openvm_recursion_circuit/transcript/merkle_verify/
air.rs

1use core::{array, borrow::Borrow};
2
3use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper};
4pub use openvm_poseidon2_air::POSEIDON2_WIDTH;
5use openvm_recursion_circuit_derive::AlignedBorrow;
6use openvm_stark_backend::{
7    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
8};
9use openvm_stark_sdk::config::baby_bear_poseidon2::{CHUNK, DIGEST_SIZE};
10use p3_air::{Air, AirBuilder, BaseAir};
11use p3_field::{Field, PrimeCharacteristicRing};
12use p3_matrix::Matrix;
13
14use crate::{
15    bus::{
16        CommitmentsBus, CommitmentsBusMessage, MerkleVerifyBus, MerkleVerifyBusMessage,
17        Poseidon2CompressBus, Poseidon2CompressMessage,
18    },
19    primitives::bus::{RightShiftBus, RightShiftMessage},
20};
21
22/// There are two parts in the merkle proof: hashing leaves and the (standard) merkle proof.
23///
24/// Example: (k = 2), going from left to right
25/// leaf0 \
26/// leaf1 -> a
27/// leaf2 \    \
28/// leaf3 -> b -> c \                 will start hashing with merkle proof siblings
29///            merkle_sibing -> ...
30///
31/// (First part) Hashing leaves: there are `2^k` leaves, each is [T; DIGEST_SIZE]. Each row in
32/// MerkleVerify AIR represents a Poseidon2 compression (of two leaves, or two intermediate values).
33/// So there are `2^k - 1` rows for this part. At height 0, the leaves are at the bottom (leaf0 ~
34/// leaf3) in the diagram above. At height 1 are the intermediate hashes (a, b) in the diagram
35/// above. The first row in the AIR will be Poseidon2 compression of leaf0 and leaf1, and the second
36/// row will be Poseidon2 compression of leaf2 and leaf3. And the third row will be Poseidon2
37/// compression of a and b.
38///
39/// (Second part) Standard merkle proof, the next row will be Poseidon2 compression of `c` and the
40/// sibling of `c`.
41#[repr(C)]
42#[derive(AlignedBorrow, StructReflection)]
43pub struct MerkleVerifyCols<T> {
44    /// Index of the proof this row is for
45    pub proof_idx: T,
46    /// Indicator: whether this row is valid
47    pub is_valid: T,
48    /// Indicator: whether this is the last row of a merkle proof
49    pub is_last_merkle: T,
50
51    /// Indicator: whether this row hashes two leaves or is part of the Merkle path proof
52    pub is_combining_leaves: T,
53    /// Indicator: whether this row is the root of the hashed leaves
54    pub is_last_leaf: T,
55    /// Row node index of this node in the leaf hash tree, 0 if is_combining_leaves is false
56    pub leaf_sub_idx: T,
57
58    /// The merkle idx of the leaf hash root
59    pub merkle_idx_bit_src: T,
60    /// Merkle idx suffix after shifting right max(0, height - k) bits
61    pub current_idx_bit_src: T,
62    /// Total depth of the merkle proof including the leaves part = merkle_proof.len() + 1 + k
63    pub total_depth: T,
64    /// 0 -> total_depth - 1, where leaves are at height 0, combined leaf hash is at height k
65    pub height: T,
66
67    pub left: [T; DIGEST_SIZE],
68    pub right: [T; DIGEST_SIZE],
69
70    /// Flag that indicates what to receive; 0 for left, 1 for right, 2 for both
71    pub recv_flag: T,
72
73    pub commit_major: T,
74    pub commit_minor: T,
75
76    pub compression_output: [T; DIGEST_SIZE],
77}
78
79#[derive(Clone, Debug)]
80pub(super) struct MerkleVerifyLog {
81    pub merkle_idx: usize,
82    pub depth: usize,
83    pub query_idx: usize,
84    /// For round 0, this is 0. For rounds > 0, this is the round index.
85    pub commit_major: usize,
86    /// For round 0, this is the commit index. For rounds > 0, this is 0.
87    pub commit_minor: usize,
88}
89
90#[derive(ColumnsAir)]
91#[columns_via(MerkleVerifyCols<u8>)]
92pub struct MerkleVerifyAir {
93    pub poseidon2_compress_bus: Poseidon2CompressBus,
94    pub merkle_verify_bus: MerkleVerifyBus,
95    pub commitments_bus: CommitmentsBus,
96    pub right_shift_bus: RightShiftBus,
97    pub k: usize,
98}
99
100impl<F: Field> BaseAir<F> for MerkleVerifyAir {
101    fn width(&self) -> usize {
102        MerkleVerifyCols::<F>::width()
103    }
104}
105
106impl<F: Field> BaseAirWithPublicValues<F> for MerkleVerifyAir {}
107impl<F: Field> PartitionedBaseAir<F> for MerkleVerifyAir {}
108
109impl<AB: AirBuilder + InteractionBuilder> Air<AB> for MerkleVerifyAir {
110    fn eval(&self, builder: &mut AB) {
111        let main = builder.main();
112        let local = main.row_slice(0).expect("window should have two elements");
113        let local: &MerkleVerifyCols<AB::Var> = (*local).borrow();
114
115        ///////////////////////////////////////////////////////////////////////
116        // Constraints
117        ///////////////////////////////////////////////////////////////////////
118        builder.assert_bool(local.is_valid);
119        builder.assert_bool(local.is_last_merkle);
120        builder.assert_bool(local.is_combining_leaves);
121        builder.assert_bool(local.is_last_leaf);
122        builder.assert_tern(local.recv_flag);
123
124        // At the last row, the cur hash should be consistent with the commit
125        for i in 0..DIGEST_SIZE {
126            builder
127                .when(local.is_last_merkle)
128                .assert_eq(local.left[i], local.right[i]);
129        }
130
131        // Last row for a Merkle tree should have height + 1 == total_depth
132        builder
133            .when(local.is_last_merkle)
134            .assert_eq(local.height + AB::Expr::ONE, local.total_depth);
135
136        // When is_combining_leaves is false, leaf_sub_idx must be 0
137        builder
138            .when(AB::Expr::ONE - local.is_combining_leaves)
139            .assert_zero(local.leaf_sub_idx);
140
141        // On the last leaf, leaf_sub_idx must be 0 and height + 1 must be k
142        builder
143            .when(local.is_last_leaf)
144            .assert_one(local.is_combining_leaves);
145        builder
146            .when(local.is_last_leaf)
147            .assert_zero(local.leaf_sub_idx);
148        builder
149            .when(local.is_last_leaf)
150            .assert_eq(local.height + AB::Expr::ONE, AB::Expr::from_usize(self.k));
151
152        // Receive both left and right for combining leaves part, otherwise receive only one
153        builder
154            .when(local.is_combining_leaves)
155            .assert_one(local.is_valid);
156        builder
157            .when(local.is_combining_leaves)
158            .assert_eq(local.recv_flag, AB::Expr::TWO);
159        builder
160            .when_ne(local.is_combining_leaves, AB::Expr::ONE)
161            .assert_bool(local.recv_flag);
162
163        // Constrain current_idx == idx on leaf rows
164        builder
165            .when(local.is_combining_leaves)
166            .assert_eq(local.merkle_idx_bit_src, local.current_idx_bit_src);
167
168        ///////////////////////////////////////////////////////////////////////
169        // Interactions
170        ///////////////////////////////////////////////////////////////////////
171
172        // This is 2x / 2x + 1 if it's a combining leaves part, otherwise it's 0.
173        let left_leaf_sub_idx = local.is_combining_leaves * local.leaf_sub_idx * AB::Expr::TWO;
174        let right_leaf_sub_idx =
175            local.is_combining_leaves * (local.leaf_sub_idx * AB::Expr::TWO + AB::Expr::ONE);
176
177        let recv_left = (local.recv_flag - AB::Expr::ONE).square();
178        let recv_right =
179            local.recv_flag * (AB::Expr::from_u8(3) - local.recv_flag) * AB::F::TWO.inverse();
180
181        let left_child_current_idx =
182            local.current_idx_bit_src * (AB::Expr::TWO - local.is_combining_leaves);
183        let right_child_current_idx =
184            left_child_current_idx.clone() + AB::Expr::ONE - local.is_combining_leaves;
185
186        self.merkle_verify_bus.receive(
187            builder,
188            local.proof_idx,
189            MerkleVerifyBusMessage {
190                value: local.left.map(Into::into),
191                merkle_idx_bit_src: local.merkle_idx_bit_src.into(),
192                current_idx_bit_src: left_child_current_idx,
193                total_depth: local.total_depth.into(),
194                height: local.height.into(),
195                is_leaf: local.is_combining_leaves.into(),
196                leaf_sub_idx: left_leaf_sub_idx,
197                commit_major: local.commit_major.into(),
198                commit_minor: local.commit_minor.into(),
199            },
200            local.is_valid * recv_left,
201        );
202        self.merkle_verify_bus.receive(
203            builder,
204            local.proof_idx,
205            MerkleVerifyBusMessage {
206                value: local.right.map(Into::into),
207                merkle_idx_bit_src: local.merkle_idx_bit_src.into(),
208                current_idx_bit_src: right_child_current_idx,
209                total_depth: local.total_depth.into(),
210                height: local.height.into(),
211                is_leaf: local.is_combining_leaves.into(),
212                leaf_sub_idx: right_leaf_sub_idx,
213                commit_major: local.commit_major.into(),
214                commit_minor: local.commit_minor.into(),
215            },
216            local.is_valid * recv_right,
217        );
218
219        self.merkle_verify_bus.send(
220            builder,
221            local.proof_idx,
222            MerkleVerifyBusMessage {
223                value: local.compression_output.map(Into::into),
224                merkle_idx_bit_src: local.merkle_idx_bit_src.into(),
225                current_idx_bit_src: local.current_idx_bit_src.into(),
226                total_depth: local.total_depth.into(),
227                height: local.height + AB::Expr::ONE,
228                is_leaf: local.is_combining_leaves - local.is_last_leaf,
229                leaf_sub_idx: local.leaf_sub_idx.into(),
230                commit_major: local.commit_major.into(),
231                commit_minor: local.commit_minor.into(),
232            },
233            (AB::Expr::ONE - local.is_last_merkle) * local.is_valid,
234        );
235
236        self.commitments_bus.lookup_key(
237            builder,
238            local.proof_idx,
239            CommitmentsBusMessage {
240                major_idx: local.commit_major,
241                minor_idx: local.commit_minor,
242                commitment: local.left, // left and right are both the commitment
243            },
244            local.is_valid * local.is_last_merkle,
245        );
246
247        self.right_shift_bus.lookup_key(
248            builder,
249            RightShiftMessage {
250                input: local.merkle_idx_bit_src.into(),
251                shift_bits: local.total_depth - AB::Expr::from_usize(self.k),
252                result: local.current_idx_bit_src.into(),
253            },
254            local.is_valid * local.is_last_merkle,
255        );
256
257        let poseidon2_input = array::from_fn(|i| {
258            if i < CHUNK {
259                local.left[i].into()
260            } else {
261                local.right[i - CHUNK].into()
262            }
263        });
264        self.poseidon2_compress_bus.lookup_key(
265            builder,
266            Poseidon2CompressMessage {
267                input: poseidon2_input,
268                output: local.compression_output.map(Into::into),
269            },
270            local.is_valid,
271        );
272    }
273}
274
275pub(super) fn compute_cum_sum<T>(values: &[T], f: impl Fn(&T) -> usize) -> Vec<usize> {
276    values
277        .iter()
278        .scan(0usize, |acc, x| {
279            *acc += f(x);
280            Some(*acc)
281        })
282        .collect()
283}
openvm_recursion_circuit/transcript/merkle_verify/air.rs

openvm_recursion_circuit/transcript/merkle_verify/
air.rs
