openvm_continuations/circuit/root/def_paths/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit::{
4    arch::instructions::DEFERRAL_AS, system::memory::dimensions::MemoryDimensions,
5};
6use openvm_circuit_primitives::{
7    utils::{and, assert_array_eq, not},
8    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
9};
10use openvm_recursion_circuit::bus::Poseidon2CompressBus;
11use openvm_recursion_circuit_derive::AlignedBorrow;
12use openvm_stark_backend::{
13    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
14};
15use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
16use p3_air::{Air, AirBuilder, BaseAir};
17use p3_field::{Field, PrimeCharacteristicRing};
18use p3_matrix::Matrix;
19
20use crate::{
21    circuit::{
22        root::bus::{
23            DeferralAccPathBus, DeferralAccPathMessage, DeferralMerkleRootsBus,
24            DeferralMerkleRootsMessage,
25        },
26        subair::{MerklePathRowView, MerklePathSubAir, MerklePathSubAirContext},
27    },
28    utils::zero_hash,
29    CommitBytes,
30};
31
32#[repr(C)]
33#[derive(AlignedBorrow, StructReflection)]
34pub struct DeferralAccMerklePathsCols<F> {
35    pub is_valid: F,
36    pub is_right_child: F,
37
38    pub initial_node_commit: [F; DIGEST_SIZE],
39    pub initial_sibling: [F; DIGEST_SIZE],
40    pub final_node_commit: [F; DIGEST_SIZE],
41    pub final_sibling: [F; DIGEST_SIZE],
42
43    pub row_idx_exp_2: F,
44    pub merkle_path_branch_bits: F,
45
46    pub depth: F,
47    pub is_skip: F,
48    pub is_within_deferral_as: F,
49    pub is_unset: F,
50}
51
52#[derive(ColumnsAir)]
53#[columns_via(DeferralAccMerklePathsCols<u8>)]
54pub struct DeferralAccMerklePathsAir {
55    pub merkle_path_subair: MerklePathSubAir,
56    pub def_acc_paths_bus: DeferralAccPathBus,
57    pub def_merkle_roots_bus: DeferralMerkleRootsBus,
58    pub address_height: usize,
59    pub zero_hash: CommitBytes,
60}
61
62impl DeferralAccMerklePathsAir {
63    pub fn new(
64        poseidon2_compress_bus: Poseidon2CompressBus,
65        def_acc_paths_bus: DeferralAccPathBus,
66        memory_merkle_roots_bus: DeferralMerkleRootsBus,
67        memory_dimensions: MemoryDimensions,
68    ) -> Self {
69        assert!(memory_dimensions.addr_space_height > 1);
70        let pv_start_idx = memory_dimensions.label_to_index((DEFERRAL_AS, 0));
71        let merkle_path_branch_bits =
72            u32::try_from(pv_start_idx).expect("merkle_path_branch_bits must fit in u32");
73        let expected_proof_len = memory_dimensions.overall_height();
74        let zero_hash = zero_hash(1).into();
75        Self {
76            merkle_path_subair: MerklePathSubAir::new(
77                poseidon2_compress_bus,
78                expected_proof_len,
79                merkle_path_branch_bits,
80            ),
81            def_acc_paths_bus,
82            def_merkle_roots_bus: memory_merkle_roots_bus,
83            address_height: memory_dimensions.address_height,
84            zero_hash,
85        }
86    }
87}
88
89impl<F> BaseAir<F> for DeferralAccMerklePathsAir {
90    fn width(&self) -> usize {
91        DeferralAccMerklePathsCols::<u8>::width()
92    }
93}
94impl<F> BaseAirWithPublicValues<F> for DeferralAccMerklePathsAir {}
95impl<F> PartitionedBaseAir<F> for DeferralAccMerklePathsAir {}
96
97impl<AB: AirBuilder + InteractionBuilder> Air<AB> for DeferralAccMerklePathsAir {
98    fn eval(&self, builder: &mut AB) {
99        let main = builder.main();
100        let (local, next) = (
101            main.row_slice(0).expect("window should have two elements"),
102            main.row_slice(1).expect("window should have two elements"),
103        );
104        let local: &DeferralAccMerklePathsCols<AB::Var> = (*local).borrow();
105        let next: &DeferralAccMerklePathsCols<AB::Var> = (*next).borrow();
106
107        /*
108         * Constrain that depth increases by 1 each row, and that all the is_skip flags
109         * are at the beginning of the trace. The first is_skip == false row is when we
110         * should receive the initial_acc_hash and final_acc_hash. If initial_acc_hash
111         * and final_acc_hash were unset, the first commit should be the compression
112         * of two zero digests.
113         */
114        builder.when_first_row().assert_zero(local.depth);
115        builder
116            .when_transition()
117            .when(local.is_valid * next.is_valid)
118            .assert_one(next.depth - local.depth);
119
120        let is_set = not(next.is_unset);
121        let is_next_start_depth = (local.is_skip - next.is_skip) * (AB::Expr::TWO - next.is_valid)
122            + local.is_unset
123                * local.is_valid
124                * (local.is_valid - AB::Expr::ONE)
125                * AB::F::TWO.inverse();
126
127        self.def_acc_paths_bus.receive(
128            builder,
129            DeferralAccPathMessage {
130                initial_acc_hash: next.initial_node_commit.map(|v| v * is_set.clone()),
131                final_acc_hash: next.final_node_commit.map(|v| v * is_set.clone()),
132                depth: next.depth * is_set,
133                is_unset: next.is_unset.into(),
134            },
135            is_next_start_depth.clone(),
136        );
137
138        assert_array_eq::<_, _, AB::Expr, _>(
139            &mut builder
140                .when(local.is_unset)
141                .when(local.is_valid)
142                .when_ne(local.is_valid, AB::F::ONE),
143            local.initial_node_commit,
144            self.zero_hash.into(),
145        );
146
147        /*
148         * Constrain that is_unset, the flag that determines if deferrals are present, is
149         * consistent over all rows.
150         */
151        builder.assert_bool(local.is_unset);
152        builder
153            .when_first_row()
154            .when(local.is_unset)
155            .assert_one(local.is_right_child);
156        builder
157            .when(and(local.is_valid, next.is_valid))
158            .assert_eq(local.is_unset, next.is_unset);
159
160        /*
161         * Constrain that is_within_deferral_as is set until address_height. We constrain
162         * the two paths to be equal as long as is_within_deferral_as is set, i.e. that the
163         * part of DEFERRAL_AS that is not included in the Merkle root is left untouched
164         * for the duration of the program execution. We also constrain that the received
165         * depth is less than or equal to the address height.
166         */
167        builder.assert_bool(local.is_within_deferral_as);
168        builder
169            .when_first_row()
170            .assert_one(local.is_within_deferral_as);
171        builder
172            .when_transition()
173            .assert_bool(local.is_within_deferral_as - next.is_within_deferral_as);
174        builder
175            .when_transition()
176            .when(local.is_within_deferral_as - next.is_within_deferral_as)
177            .assert_eq(next.depth, AB::Expr::from_usize(self.address_height));
178        builder
179            .when_last_row()
180            .assert_zero(local.is_within_deferral_as);
181        builder
182            .when(not(local.is_within_deferral_as))
183            .assert_zero(is_next_start_depth);
184
185        assert_array_eq(
186            &mut builder.when(local.is_within_deferral_as),
187            local.initial_sibling,
188            local.final_sibling,
189        );
190
191        assert_array_eq(
192            &mut builder.when(and(local.is_within_deferral_as, local.is_unset)),
193            local.initial_node_commit,
194            local.final_node_commit,
195        );
196
197        /*
198         * Call the Merkle path sub-AIR on both the initial and final paths
199         */
200        self.merkle_path_subair.eval(
201            builder,
202            (
203                MerklePathSubAirContext {
204                    local: MerklePathRowView {
205                        is_valid: &local.is_valid,
206                        is_right_child: &local.is_right_child,
207                        node_commit: &local.initial_node_commit,
208                        sibling: &local.initial_sibling,
209                        row_idx_exp_2: &local.row_idx_exp_2,
210                        merkle_path_branch_bits: &local.merkle_path_branch_bits,
211                    },
212                    next: MerklePathRowView {
213                        is_valid: &next.is_valid,
214                        is_right_child: &next.is_right_child,
215                        node_commit: &next.initial_node_commit,
216                        sibling: &next.initial_sibling,
217                        row_idx_exp_2: &next.row_idx_exp_2,
218                        merkle_path_branch_bits: &next.merkle_path_branch_bits,
219                    },
220                },
221                local.is_skip.into(),
222                next.is_skip.into(),
223                local.is_unset.into(),
224            ),
225        );
226
227        self.merkle_path_subair.eval(
228            builder,
229            (
230                MerklePathSubAirContext {
231                    local: MerklePathRowView {
232                        is_valid: &local.is_valid,
233                        is_right_child: &local.is_right_child,
234                        node_commit: &local.final_node_commit,
235                        sibling: &local.final_sibling,
236                        row_idx_exp_2: &local.row_idx_exp_2,
237                        merkle_path_branch_bits: &local.merkle_path_branch_bits,
238                    },
239                    next: MerklePathRowView {
240                        is_valid: &next.is_valid,
241                        is_right_child: &next.is_right_child,
242                        node_commit: &next.final_node_commit,
243                        sibling: &next.final_sibling,
244                        row_idx_exp_2: &next.row_idx_exp_2,
245                        merkle_path_branch_bits: &next.merkle_path_branch_bits,
246                    },
247                },
248                local.is_skip.into(),
249                next.is_skip.into(),
250                local.is_unset.into(),
251            ),
252        );
253
254        /*
255         * Receive the final memory merkle root on the last valid row.
256         */
257        self.def_merkle_roots_bus.receive(
258            builder,
259            DeferralMerkleRootsMessage {
260                initial_root: local.initial_node_commit,
261                final_root: local.final_node_commit,
262            },
263            local.is_valid * (AB::Expr::ONE - next.is_valid).square(),
264        );
265    }
266}
openvm_continuations/circuit/root/def_paths/air.rs

openvm_continuations/circuit/root/def_paths/
air.rs
