openvm_circuit/system/memory/merkle/
air.rs

1use std::{borrow::Borrow, iter};
2
3use openvm_stark_backend::{
4    interaction::{InteractionBuilder, PermutationCheckBus},
5    p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir},
6    p3_field::{Field, PrimeCharacteristicRing},
7    p3_matrix::Matrix,
8    rap::{BaseAirWithPublicValues, PartitionedBaseAir},
9};
10
11use crate::system::memory::merkle::{MemoryDimensions, MemoryMerkleCols, MemoryMerklePvs};
12
13#[derive(Clone, Debug)]
14pub struct MemoryMerkleAir<const CHUNK: usize> {
15    pub memory_dimensions: MemoryDimensions,
16    pub merkle_bus: PermutationCheckBus,
17    pub compression_bus: PermutationCheckBus,
18}
19
20impl<const CHUNK: usize, F: Field> PartitionedBaseAir<F> for MemoryMerkleAir<CHUNK> {}
21impl<const CHUNK: usize, F: Field> BaseAir<F> for MemoryMerkleAir<CHUNK> {
22    fn width(&self) -> usize {
23        MemoryMerkleCols::<F, CHUNK>::width()
24    }
25}
26impl<const CHUNK: usize, F: Field> BaseAirWithPublicValues<F> for MemoryMerkleAir<CHUNK> {
27    fn num_public_values(&self) -> usize {
28        MemoryMerklePvs::<F, CHUNK>::width()
29    }
30}
31
32impl<const CHUNK: usize, AB: InteractionBuilder + AirBuilderWithPublicValues> Air<AB>
33    for MemoryMerkleAir<CHUNK>
34{
35    fn eval(&self, builder: &mut AB) {
36        let main = builder.main();
37        let (local, next) = (
38            main.row_slice(0).expect("window should have two elements"),
39            main.row_slice(1).expect("window should have two elements"),
40        );
41        let local: &MemoryMerkleCols<_, CHUNK> = (*local).borrow();
42        let next: &MemoryMerkleCols<_, CHUNK> = (*next).borrow();
43
44        // `expand_direction` should be -1, 0, 1
45        builder.assert_eq(
46            local.expand_direction,
47            local.expand_direction * local.expand_direction * local.expand_direction,
48        );
49
50        builder.assert_bool(local.left_direction_different);
51        builder.assert_bool(local.right_direction_different);
52
53        // if `expand_direction` != -1, then `*_direction_different` should be 0
54        builder
55            .when_ne(local.expand_direction, AB::F::NEG_ONE)
56            .assert_zero(local.left_direction_different);
57        builder
58            .when_ne(local.expand_direction, AB::F::NEG_ONE)
59            .assert_zero(local.right_direction_different);
60
61        // rows should be sorted in descending order
62        // independently by `parent_height`, `height_section`, `is_root`
63        builder
64            .when_transition()
65            .assert_bool(local.parent_height - next.parent_height);
66        builder
67            .when_transition()
68            .assert_bool(local.height_section - next.height_section);
69        builder
70            .when_transition()
71            .assert_bool(local.is_root - next.is_root);
72
73        // row with greatest height should have `height_section` = 1
74        builder.when_first_row().assert_one(local.height_section);
75        // two rows with greatest height should have `is_root` = 1
76        builder.when_first_row().assert_one(local.is_root);
77        builder.when_first_row().assert_one(next.is_root);
78        // the root rows have `as_label` and `address_label` set to zero,
79        // so that we can't use another tree representation
80        builder
81            .when(local.is_root)
82            .assert_zero(local.parent_address_label);
83        builder
84            .when(local.is_root)
85            .assert_zero(local.parent_as_label);
86        // row with least height should have `height_section` = 0, `is_root` = 0
87        builder.when_last_row().assert_zero(local.height_section);
88        builder.when_last_row().assert_zero(local.is_root);
89        // `height_section` changes from 0 to 1 only when `parent_height` changes from
90        // `address_height` to `address_height` + 1
91        builder
92            .when_transition()
93            .when_ne(
94                local.parent_height,
95                AB::F::from_usize(self.memory_dimensions.address_height + 1),
96            )
97            .assert_eq(local.height_section, next.height_section);
98        builder
99            .when_transition()
100            .when_ne(
101                next.parent_height,
102                AB::F::from_usize(self.memory_dimensions.address_height),
103            )
104            .assert_eq(local.height_section, next.height_section);
105        // two adjacent rows with `is_root` = 1 should have
106        // the first `expand_direction` = 1, the second `expand_direction` = -1
107        builder
108            .when(local.is_root)
109            .when(next.is_root)
110            .assert_eq(local.expand_direction - next.expand_direction, AB::F::TWO);
111
112        // roots should have correct height
113        builder.when(local.is_root).assert_eq(
114            local.parent_height,
115            AB::Expr::from_usize(self.memory_dimensions.overall_height()),
116        );
117
118        // constrain public values
119        let &MemoryMerklePvs::<_, CHUNK> {
120            initial_root,
121            final_root,
122        } = builder.public_values().borrow();
123        for i in 0..CHUNK {
124            builder
125                .when_first_row()
126                .assert_eq(local.parent_hash[i], initial_root[i]);
127            builder
128                .when_first_row()
129                .assert_eq(next.parent_hash[i], final_root[i]);
130        }
131
132        self.eval_interactions(builder, local);
133    }
134}
135
136impl<const CHUNK: usize> MemoryMerkleAir<CHUNK> {
137    pub fn eval_interactions<AB: InteractionBuilder>(
138        &self,
139        builder: &mut AB,
140        local: &MemoryMerkleCols<AB::Var, CHUNK>,
141    ) {
142        // interaction does not occur for first two rows;
143        // for those, parent hash value comes from public values
144        self.merkle_bus.interact(
145            builder,
146            [
147                local.expand_direction.into(),
148                local.parent_height.into(),
149                local.parent_as_label.into(),
150                local.parent_address_label.into(),
151            ]
152            .into_iter()
153            .chain(local.parent_hash.into_iter().map(Into::into)),
154            // count can probably be made degree 1 if necessary
155            (AB::Expr::ONE - local.is_root) * local.expand_direction,
156        );
157
158        self.merkle_bus.interact(
159            builder,
160            [
161                local.expand_direction + (local.left_direction_different * AB::F::TWO),
162                local.parent_height - AB::F::ONE,
163                local.parent_as_label * (AB::Expr::ONE + local.height_section),
164                local.parent_address_label * (AB::Expr::TWO - local.height_section),
165            ]
166            .into_iter()
167            .chain(local.left_child_hash.into_iter().map(Into::into)),
168            -local.expand_direction.into(),
169        );
170
171        self.merkle_bus.interact(
172            builder,
173            [
174                local.expand_direction + (local.right_direction_different * AB::F::TWO),
175                local.parent_height - AB::F::ONE,
176                (local.parent_as_label * (AB::Expr::ONE + local.height_section))
177                    + local.height_section,
178                (local.parent_address_label * (AB::Expr::TWO - local.height_section))
179                    + (AB::Expr::ONE - local.height_section),
180            ]
181            .into_iter()
182            .chain(local.right_child_hash.into_iter().map(Into::into)),
183            -local.expand_direction.into(),
184        );
185
186        let compress_fields = iter::empty()
187            .chain(local.left_child_hash)
188            .chain(local.right_child_hash)
189            .chain(local.parent_hash);
190        self.compression_bus.interact(
191            builder,
192            compress_fields,
193            local.expand_direction * local.expand_direction,
194        );
195    }
196}
openvm_circuit/system/memory/merkle/air.rs

openvm_circuit/system/memory/merkle/
air.rs
