openvm_continuations/circuit/subair/merkle_path/
air.rs

1use std::array::from_fn;
2
3use itertools::Itertools;
4use openvm_circuit::arch::POSEIDON2_WIDTH;
5use openvm_circuit_primitives::{utils::not, SubAir};
6use openvm_recursion_circuit::bus::{Poseidon2CompressBus, Poseidon2CompressMessage};
7use openvm_stark_backend::interaction::InteractionBuilder;
8use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
9use p3_air::AirBuilder;
10use p3_field::PrimeCharacteristicRing;
11
12#[derive(Copy, Clone, Debug)]
13pub struct MerklePathRowView<'a, T> {
14    // 0 for invalid, 1 for valid, 2 for valid first row
15    pub is_valid: &'a T,
16    pub is_right_child: &'a T,
17    pub node_commit: &'a [T; DIGEST_SIZE],
18    pub sibling: &'a [T; DIGEST_SIZE],
19    // 2^row_idx, used to (a) constrain merkle proof height and (b) accumulate
20    // merkle_path_branch_bits
21    pub row_idx_exp_2: &'a T,
22    pub merkle_path_branch_bits: &'a T,
23}
24
25#[derive(Copy, Clone, Debug)]
26pub struct MerklePathSubAirContext<'a, T> {
27    pub local: MerklePathRowView<'a, T>,
28    pub next: MerklePathRowView<'a, T>,
29}
30
31/// SubAir to constrain a single Merkle path represented over consecutive rows.
32#[derive(Clone, Debug, derive_new::new)]
33pub struct MerklePathSubAir {
34    pub poseidon2_bus: Poseidon2CompressBus,
35    pub expected_proof_len: usize,
36    pub expected_branch_bits: u32,
37}
38
39impl<AB: AirBuilder + InteractionBuilder> SubAir<AB> for MerklePathSubAir {
40    type AirContext<'a>
41        = (
42        MerklePathSubAirContext<'a, AB::Var>,
43        AB::Expr,
44        AB::Expr,
45        AB::Expr,
46    )
47    where
48        AB: 'a,
49        AB::Var: 'a,
50        AB::Expr: 'a;
51
52    fn eval<'a>(&'a self, builder: &'a mut AB, ctx: Self::AirContext<'a>)
53    where
54        AB::Var: 'a,
55        AB::Expr: 'a,
56    {
57        let local = ctx.0.local;
58        let next = ctx.0.next;
59
60        builder.assert_tern(*local.is_valid);
61        builder.assert_bool(*local.is_right_child);
62
63        /*
64         * Constrain that is_valid is 2 on the first row, that the rest of the Merkle
65         * tree immediately follows, and that all invalid rows are at the end.
66         */
67        builder
68            .when_first_row()
69            .assert_eq(*local.is_valid, AB::F::TWO);
70        builder
71            .when_transition()
72            .assert_bool(*local.is_valid - *next.is_valid);
73        builder.when_transition().assert_bool(*next.is_valid);
74
75        /*
76         * Constrain that the Merkle path starts at depth 0 (row_idx_exp_2 = 1), then
77         * doubles every valid row. The accumulated branch bits start from is_right_child
78         * on the first row and must end at expected_branch_bits.
79         */
80        builder.when_first_row().assert_one(*local.row_idx_exp_2);
81        builder
82            .when_first_row()
83            .assert_eq(*local.merkle_path_branch_bits, *local.is_right_child);
84
85        let mut when_transition = builder.when_transition();
86        let mut when_both_valid = when_transition.when(*next.is_valid);
87        when_both_valid.assert_eq(*local.row_idx_exp_2 * AB::F::TWO, *next.row_idx_exp_2);
88        when_both_valid.assert_eq(
89            *local.merkle_path_branch_bits + *next.is_right_child * *next.row_idx_exp_2,
90            *next.merkle_path_branch_bits,
91        );
92
93        let expected_branch_bits_offset = ctx.3;
94        let mut when_last = builder.when(*local.is_valid * (*next.is_valid - AB::F::ONE));
95        when_last.assert_eq(
96            *local.merkle_path_branch_bits,
97            expected_branch_bits_offset + AB::F::from_u32(self.expected_branch_bits),
98        );
99        when_last.assert_eq(
100            *local.row_idx_exp_2,
101            AB::F::from_usize(1 << self.expected_proof_len),
102        );
103        when_last.assert_zero(*local.is_right_child);
104
105        /*
106         * We give the option to skip hashing until some row_idx, effectively allowing
107         * the Merkle proof to begin there. The caller MUST constrain the validity of
108         * node_commit and row_idx on that row themselves. Additionally, there are some
109         * restrictions:
110         * - All skips must happen at the beginning
111         * - is_valid  must be set as normal
112         * - is_right_child and merkle_path_branch_bits must be 0
113         */
114        let local_skip_hash = ctx.1;
115        let next_skip_hash = ctx.2;
116
117        builder.assert_bool(local_skip_hash.clone());
118        builder
119            .when_transition()
120            .assert_bool(local_skip_hash.clone() - next_skip_hash.clone());
121        builder
122            .when_transition()
123            .when(local_skip_hash.clone() - next_skip_hash)
124            .assert_one(*next.is_valid);
125
126        let mut when_skip_hash = builder.when(local_skip_hash.clone());
127        when_skip_hash.assert_zero(*local.is_right_child);
128        when_skip_hash.assert_zero(*local.merkle_path_branch_bits);
129
130        /*
131         * Constrain that next.node_commit is the Poseidon2 compression of local_commit
132         * and sibling. Note that node_commit on the first row is unconstrained.
133         */
134        let left: [AB::Expr; DIGEST_SIZE] = from_fn(|i| {
135            (AB::Expr::ONE - *local.is_right_child) * local.node_commit[i]
136                + *local.is_right_child * local.sibling[i]
137        });
138        let right: [AB::Expr; DIGEST_SIZE] = from_fn(|i| {
139            *local.is_right_child * local.node_commit[i]
140                + (AB::Expr::ONE - *local.is_right_child) * local.sibling[i]
141        });
142
143        let poseidon2_input: [AB::Expr; POSEIDON2_WIDTH] =
144            left.into_iter().chain(right).collect_array().unwrap();
145        let should_hash = *next.is_valid * (AB::Expr::TWO - *next.is_valid);
146
147        self.poseidon2_bus.lookup_key(
148            builder,
149            Poseidon2CompressMessage {
150                input: poseidon2_input,
151                output: next.node_commit.map(Into::into),
152            },
153            should_hash * not::<AB::Expr>(local_skip_hash),
154        );
155    }
156}
openvm_continuations/circuit/subair/merkle_path/air.rs

openvm_continuations/circuit/subair/merkle_path/
air.rs
