openvm_continuations/circuit/subair/merkle_tree/
air.rs

1use std::array::from_fn;
2
3use openvm_circuit_primitives::{StructReflection, StructReflectionHelper, SubAir};
4use openvm_recursion_circuit::bus::{Poseidon2CompressBus, Poseidon2CompressMessage};
5use openvm_recursion_circuit_derive::AlignedBorrow;
6use openvm_stark_backend::interaction::InteractionBuilder;
7use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
8use p3_air::AirBuilder;
9use p3_field::{Field, PrimeCharacteristicRing};
10
11use crate::{
12    circuit::{
13        deferral::{DEF_INTERNAL_TAG, DEF_LEAF_TAG},
14        subair::merkle_tree::bus::{
15            MerkleRootBus, MerkleRootMessage, MerkleTreeInternalBus, MerkleTreeInternalMessage,
16        },
17    },
18    utils::digests_to_poseidon2_input,
19};
20
21#[repr(C)]
22#[derive(AlignedBorrow, StructReflection, Copy, Clone, Debug)]
23pub struct MerkleTreeCols<T> {
24    pub row_idx: T,
25
26    // 1 for internal, 2 for root, 0 for invalid
27    pub send_type: T,
28    // 1 for leaf, 2 for internal, 0 for invalid
29    pub receive_type: T,
30
31    pub parent: [T; DIGEST_SIZE],
32    pub is_right_child: T,
33
34    pub left_child: [T; DIGEST_SIZE],
35    pub right_child: [T; DIGEST_SIZE],
36}
37
38/// SubAir to constrain a Merkle root proof and send the root to MerkleRootBus
39#[derive(Clone, Debug, derive_new::new)]
40pub struct MerkleTreeSubAir {
41    pub poseidon2_bus: Poseidon2CompressBus,
42    pub merkle_root_bus: MerkleRootBus,
43    pub merkle_tree_internal_bus: MerkleTreeInternalBus,
44
45    // Identifier in case multiple AIRs use this sub-AIR
46    pub idx: usize,
47    // Flag that indicates whether we should send row_idx (i.e. num_nodes) on
48    // the root row or not
49    pub send_num_rows: bool,
50}
51
52impl<AB: AirBuilder + InteractionBuilder> SubAir<AB> for MerkleTreeSubAir {
53    type AirContext<'a>
54        = (
55        &'a MerkleTreeCols<AB::Var>,
56        &'a MerkleTreeCols<AB::Var>,
57        AB::Expr,
58        Option<&'a [AB::Var; DIGEST_SIZE]>,
59    )
60    where
61        AB: 'a,
62        AB::Var: 'a,
63        AB::Expr: 'a;
64
65    fn eval<'a>(&'a self, builder: &'a mut AB, ctx: Self::AirContext<'a>)
66    where
67        AB::Var: 'a,
68        AB::Expr: 'a,
69    {
70        let (local, next, num_rows, tagged_left_child) = ctx;
71
72        /*
73         * Constrain that row_idx actually corresponds to the matrix row index, and
74         * that there are exactly num_rows rows.
75         */
76        builder.when_first_row().assert_zero(local.row_idx);
77        builder
78            .when_transition()
79            .assert_eq(local.row_idx + AB::F::ONE, next.row_idx);
80        builder
81            .when_last_row()
82            .assert_eq(local.row_idx, num_rows.clone() - AB::F::ONE);
83
84        /*
85         * Constrain that is_right_child is correctly set. Even row_idx correspond to
86         * left children, and odd row_idx correspond to right - thus we constrain that
87         * the is_right_child flag alternates.
88         */
89        builder.when_first_row().assert_zero(local.is_right_child);
90        builder.assert_eq(local.is_right_child, AB::Expr::ONE - next.is_right_child);
91
92        /*
93         * Constrain that the first num_rows - 2 rows send their parent internally,
94         * the second last row sends the commit, and the last row sends nothing. There
95         * must be at least 2 rows, which is enforced by constraining send_type to be
96         * 1 or 2 on the first row and 0 on the last.
97         */
98        builder.assert_tern(local.send_type);
99
100        builder
101            .when_first_row()
102            .assert_bool(local.send_type - AB::F::ONE);
103        builder
104            .when_transition()
105            .assert_bool(local.send_type - AB::F::ONE);
106        builder.when_last_row().assert_zero(local.send_type);
107        builder
108            .when_ne(local.send_type, AB::F::TWO)
109            .assert_bool(next.send_type - AB::F::ONE);
110        builder
111            .when_ne(local.send_type, AB::F::ZERO)
112            .when_ne(local.send_type, AB::F::ONE)
113            .assert_zero(next.send_type);
114
115        /*
116         * Constrain that the first num_rows / 2 rows read their left and right values
117         * from pvs, the next (num_rows / 2) - 1 rows receive their values internally,
118         * and that the last row receives no values. We constrain this by asserting the
119         * last row's receive_type == 0 and that elsewhere the column is monotonically
120         * increasing.
121         */
122        builder.assert_tern(local.receive_type);
123
124        builder.when_first_row().assert_one(local.receive_type);
125        builder
126            .when_transition()
127            .assert_bool(local.receive_type - AB::F::ONE);
128        builder.when_last_row().assert_zero(local.receive_type);
129
130        builder
131            .when_ne(local.receive_type, AB::F::ZERO)
132            .when_ne(local.receive_type, AB::F::ONE)
133            .assert_zero(next.receive_type * (next.receive_type - AB::F::TWO));
134
135        /*
136         * Send and receive internal messages. Given a non-root node at index i, its
137         * parent is at index floor((i + num_rows) / 2). Note that we send the
138         * internal message when local.send_type == 1, and receive left and right
139         * when internal_receive == 2.
140         */
141        let half = AB::F::TWO.inverse();
142        let internal_send = local.send_type * (AB::Expr::TWO - local.send_type);
143        let internal_receive = local.receive_type * (local.receive_type - AB::F::ONE) * half;
144
145        self.merkle_tree_internal_bus.send(
146            builder,
147            MerkleTreeInternalMessage {
148                child_value: local.parent.map(Into::into),
149                is_right_child: local.is_right_child.into(),
150                parent_idx: (local.row_idx - local.is_right_child + num_rows.clone()) * half,
151            },
152            internal_send,
153        );
154
155        self.merkle_tree_internal_bus.receive(
156            builder,
157            MerkleTreeInternalMessage {
158                child_value: local.left_child.map(Into::into),
159                is_right_child: AB::Expr::ZERO,
160                parent_idx: local.row_idx.into(),
161            },
162            internal_receive.clone(),
163        );
164
165        self.merkle_tree_internal_bus.receive(
166            builder,
167            MerkleTreeInternalMessage {
168                child_value: local.right_child.map(Into::into),
169                is_right_child: AB::Expr::ONE,
170                parent_idx: local.row_idx.into(),
171            },
172            internal_receive,
173        );
174
175        /*
176         * Send and receive external messages. At every valid node we have to receive a
177         * message from the Poseidon2 peripheral AIR to constrain that parent is the
178         * compression hash of left_child and right_child. At the root node we send the
179         * commit to MerkleRootBus to be used elsehwhere. Note a row is valid iff it is
180         * not the last, i.e. iff send_type (or equivalently receive_type) is non-zero.
181         */
182        let is_valid = local.send_type * (AB::Expr::from_u8(3) - local.send_type) * half;
183        let is_root = local.send_type * (local.send_type - AB::F::ONE) * half;
184
185        if let Some(tagged_left_child) = tagged_left_child {
186            let is_leaf = local.receive_type * (AB::Expr::TWO - local.receive_type);
187            let is_internal = local.receive_type * (local.receive_type - AB::F::ONE) * half;
188            let tag = from_fn(|i| {
189                is_leaf.clone() * AB::Expr::from_u8(DEF_LEAF_TAG[i])
190                    + is_internal.clone() * AB::Expr::from_u8(DEF_INTERNAL_TAG[i])
191            });
192
193            self.poseidon2_bus.lookup_key(
194                builder,
195                Poseidon2CompressMessage {
196                    input: digests_to_poseidon2_input(tag, local.left_child.map(Into::into)),
197                    output: tagged_left_child.map(Into::into),
198                },
199                is_valid.clone(),
200            );
201            self.poseidon2_bus.lookup_key(
202                builder,
203                Poseidon2CompressMessage {
204                    input: digests_to_poseidon2_input(
205                        tagged_left_child.map(Into::into),
206                        local.right_child.map(Into::into),
207                    ),
208                    output: local.parent.map(Into::into),
209                },
210                is_valid,
211            );
212        } else {
213            self.poseidon2_bus.lookup_key(
214                builder,
215                Poseidon2CompressMessage {
216                    input: digests_to_poseidon2_input(local.left_child, local.right_child),
217                    output: local.parent,
218                },
219                is_valid,
220            );
221        }
222
223        self.merkle_root_bus.send(
224            builder,
225            MerkleRootMessage {
226                merkle_root: local.parent.map(Into::into),
227                idx: AB::Expr::from_usize(self.idx),
228                num_rows_or_zero: if self.send_num_rows {
229                    num_rows
230                } else {
231                    AB::Expr::ZERO
232                },
233            },
234            is_root,
235        );
236    }
237}
openvm_continuations/circuit/subair/merkle_tree/air.rs

openvm_continuations/circuit/subair/merkle_tree/
air.rs
