openvm_circuit/system/memory/merkle/
trace.rs

1use std::{borrow::BorrowMut, sync::atomic::AtomicU32};
2
3use openvm_cpu_backend::CpuBackend;
4use openvm_stark_backend::{
5    p3_field::PrimeField32, p3_matrix::dense::RowMajorMatrix, prover::AirProvingContext,
6    StarkProtocolConfig, Val,
7};
8use tracing::instrument;
9
10use crate::{
11    arch::{hasher::HasherChip, VmField},
12    system::{
13        memory::{
14            merkle::{tree::MerkleTree, FinalState, MemoryMerkleChip, MemoryMerkleCols},
15            Equipartition, MemoryImage,
16        },
17        poseidon2::{
18            Poseidon2PeripheryBaseChip, Poseidon2PeripheryChip, PERIPHERY_POSEIDON2_WIDTH,
19        },
20    },
21};
22
23impl<const CHUNK: usize, F: PrimeField32> MemoryMerkleChip<CHUNK, F> {
24    #[instrument(name = "merkle_finalize", level = "debug", skip_all)]
25    pub(crate) fn finalize(
26        &mut self,
27        initial_memory: &MemoryImage,
28        final_memory: &Equipartition<F, CHUNK>,
29        hasher: &impl HasherChip<CHUNK, F>,
30    ) {
31        assert!(self.final_state.is_none(), "Merkle chip already finalized");
32        let memory_dimensions = &self.air.memory_dimensions;
33        let mut tree = MerkleTree::from_memory(initial_memory, memory_dimensions, hasher);
34        self.final_state = Some(tree.finalize(hasher, final_memory, memory_dimensions));
35        self.top_tree = tree.top_tree(memory_dimensions.addr_space_height);
36    }
37}
38
39impl<const CHUNK: usize, F> MemoryMerkleChip<CHUNK, F>
40where
41    F: PrimeField32,
42{
43    pub fn generate_proving_ctx<SC>(&mut self) -> AirProvingContext<CpuBackend<SC>>
44    where
45        SC: StarkProtocolConfig<F = F>,
46    {
47        assert!(
48            self.final_state.is_some(),
49            "Merkle chip must finalize before trace generation"
50        );
51        let FinalState {
52            mut rows,
53            init_root,
54            final_root,
55        } = self.final_state.take().unwrap();
56        // important that this sort be stable,
57        // because we need the initial root to be first and the final root to be second
58        rows.reverse();
59        rows.swap(0, 1);
60
61        #[cfg(feature = "metrics")]
62        {
63            self.current_height = rows.len();
64        }
65        let width = MemoryMerkleCols::<Val<SC>, CHUNK>::width();
66        let mut height = rows.len().next_power_of_two();
67        if let Some(mut oh) = self.overridden_height {
68            oh = oh.next_power_of_two();
69            assert!(
70                oh >= height,
71                "Overridden height {oh} is less than the required height {height}"
72            );
73            height = oh;
74        }
75        let mut trace = Val::<SC>::zero_vec(width * height);
76
77        for (trace_row, row) in trace.chunks_exact_mut(width).zip(rows) {
78            *trace_row.borrow_mut() = row;
79        }
80
81        let trace = RowMajorMatrix::new(trace, width);
82        let pvs = init_root.into_iter().chain(final_root).collect();
83        AirProvingContext::simple(trace, pvs)
84    }
85}
86pub trait SerialReceiver<T> {
87    fn receive(&self, msg: T);
88}
89
90impl<'a, F: VmField, const SBOX_REGISTERS: usize> SerialReceiver<&'a [F]>
91    for Poseidon2PeripheryBaseChip<F, SBOX_REGISTERS>
92{
93    /// Receives a permutation preimage, pads with zeros to the permutation width, and records.
94    /// The permutation preimage must have length at most the permutation width (panics otherwise).
95    fn receive(&self, perm_preimage: &'a [F]) {
96        assert!(perm_preimage.len() <= PERIPHERY_POSEIDON2_WIDTH);
97        let mut state = [F::ZERO; PERIPHERY_POSEIDON2_WIDTH];
98        state[..perm_preimage.len()].copy_from_slice(perm_preimage);
99        let count = self.records.entry(state).or_insert(AtomicU32::new(0));
100        count.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
101        self.nonempty
102            .store(true, std::sync::atomic::Ordering::Relaxed);
103    }
104}
105
106impl<'a, F: VmField> SerialReceiver<&'a [F]> for Poseidon2PeripheryChip<F> {
107    fn receive(&self, perm_preimage: &'a [F]) {
108        match self {
109            Poseidon2PeripheryChip::Register0(chip) => chip.receive(perm_preimage),
110            Poseidon2PeripheryChip::Register1(chip) => chip.receive(perm_preimage),
111        }
112    }
113}
openvm_circuit/system/memory/merkle/trace.rs

openvm_circuit/system/memory/merkle/
trace.rs
