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    /// `compression_bus` is the bus for direct (no-memory involved) interactions to call the cryptographic compression function.
32    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}