openvm_circuit/system/memory/merkle/
mod.rs
1use openvm_stark_backend::{interaction::PermutationCheckBus, p3_field::PrimeField32};
2use rustc_hash::FxHashSet;
3
4use super::controller::dimensions::MemoryDimensions;
5mod air;
6mod columns;
7mod trace;
8
9pub use air::*;
10pub use columns::*;
11pub(super) use trace::SerialReceiver;
12
13#[cfg(test)]
14mod tests;
15
16pub struct MemoryMerkleChip<const CHUNK: usize, F> {
17 pub air: MemoryMerkleAir<CHUNK>,
18 touched_nodes: FxHashSet<(usize, u32, u32)>,
19 num_touched_nonleaves: usize,
20 final_state: Option<FinalState<CHUNK, F>>,
21 overridden_height: Option<usize>,
22}
23#[derive(Debug)]
24struct FinalState<const CHUNK: usize, F> {
25 rows: Vec<MemoryMerkleCols<F, CHUNK>>,
26 init_root: [F; CHUNK],
27 final_root: [F; CHUNK],
28}
29
30impl<const CHUNK: usize, F: PrimeField32> MemoryMerkleChip<CHUNK, F> {
31 pub fn new(
33 memory_dimensions: MemoryDimensions,
34 merkle_bus: PermutationCheckBus,
35 compression_bus: PermutationCheckBus,
36 ) -> Self {
37 assert!(memory_dimensions.as_height > 0);
38 assert!(memory_dimensions.address_height > 0);
39 let mut touched_nodes = FxHashSet::default();
40 touched_nodes.insert((memory_dimensions.overall_height(), 0, 0));
41 Self {
42 air: MemoryMerkleAir {
43 memory_dimensions,
44 merkle_bus,
45 compression_bus,
46 },
47 touched_nodes,
48 num_touched_nonleaves: 1,
49 final_state: None,
50 overridden_height: None,
51 }
52 }
53 pub fn set_overridden_height(&mut self, override_height: usize) {
54 self.overridden_height = Some(override_height);
55 }
56
57 fn touch_node(&mut self, height: usize, as_label: u32, address_label: u32) {
58 if self.touched_nodes.insert((height, as_label, address_label)) {
59 assert_ne!(height, self.air.memory_dimensions.overall_height());
60 if height != 0 {
61 self.num_touched_nonleaves += 1;
62 }
63 if height >= self.air.memory_dimensions.address_height {
64 self.touch_node(height + 1, as_label / 2, address_label);
65 } else {
66 self.touch_node(height + 1, as_label, address_label / 2);
67 }
68 }
69 }
70
71 pub fn touch_range(&mut self, address_space: u32, address: u32, len: u32) {
72 let as_label = address_space - self.air.memory_dimensions.as_offset;
73 let first_address_label = address / CHUNK as u32;
74 let last_address_label = (address + len - 1) / CHUNK as u32;
75 for address_label in first_address_label..=last_address_label {
76 self.touch_node(0, as_label, address_label);
77 }
78 }
79}