openvm_circuit/system/memory/merkle/
air.rs

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

openvm_circuit/system/memory/merkle/
air.rs
