openvm_circuit/system/memory/
persistent.rs

1use std::{
2    array,
3    borrow::{Borrow, BorrowMut},
4    iter,
5};
6
7use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper};
8use openvm_circuit_primitives_derive::AlignedBorrow;
9use openvm_cpu_backend::CpuBackend;
10use openvm_stark_backend::{
11    interaction::{InteractionBuilder, PermutationCheckBus},
12    p3_air::{Air, AirBuilder, BaseAir},
13    p3_field::{PrimeCharacteristicRing, PrimeField32},
14    p3_matrix::{dense::RowMajorMatrix, Matrix},
15    p3_maybe_rayon::prelude::*,
16    prover::AirProvingContext,
17    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
18};
19use tracing::instrument;
20
21use super::{merkle::SerialReceiver, online::INITIAL_TIMESTAMP};
22use crate::{
23    arch::{hasher::Hasher, ADDR_SPACE_OFFSET, DEFAULT_BLOCK_SIZE},
24    primitives::Chip,
25    system::memory::{
26        controller::CHUNK, offline_checker::MemoryBus, MemoryAddress, MemoryImage,
27        TimestampedEquipartition,
28    },
29};
30
31/// Number of DEFAULT_BLOCK_SIZE blocks per CHUNK (e.g., 2 for 8/4).
32/// Blocks are on the same row only for Merkle tree hashing (8 bytes at a time).
33/// Memory bus interactions use per-block timestamps.
34pub const BLOCKS_PER_CHUNK: usize = CHUNK / DEFAULT_BLOCK_SIZE;
35
36/// The values describe aligned chunk of memory of size `CHUNK`---the data together with the last
37/// accessed timestamp---in either the initial or final memory state.
38#[repr(C)]
39#[derive(Debug, AlignedBorrow, StructReflection)]
40pub struct PersistentBoundaryCols<T, const CHUNK: usize> {
41    // `expand_direction` =  1 corresponds to initial memory state
42    // `expand_direction` = -1 corresponds to final memory state
43    // `expand_direction` =  0 corresponds to irrelevant row (all interactions multiplicity 0)
44    pub expand_direction: T,
45    pub address_space: T,
46    pub leaf_label: T,
47    pub values: [T; CHUNK],
48    pub hash: [T; CHUNK],
49    /// Per-block timestamps. Each DEFAULT_BLOCK_SIZE block within the chunk has its own timestamp.
50    /// For untouched blocks, timestamp stays at 0 (balances: boundary sends at t=0 init, receives
51    /// at t=0 final).
52    pub timestamps: [T; BLOCKS_PER_CHUNK],
53}
54
55/// Imposes the following constraints:
56/// - `expand_direction` should be -1, 0, 1
57///
58/// Sends the following interactions:
59/// - if `expand_direction` is 1, sends `[0, 0, address_space_label, leaf_label]` to `merkle_bus`.
60/// - if `expand_direction` is -1, receives `[1, 0, address_space_label, leaf_label]` from
61///   `merkle_bus`.
62#[derive(Clone, Debug, ColumnsAir)]
63#[columns_via(PersistentBoundaryCols<u8, CHUNK>)]
64pub struct PersistentBoundaryAir<const CHUNK: usize> {
65    pub memory_bus: MemoryBus,
66    pub merkle_bus: PermutationCheckBus,
67    pub compression_bus: PermutationCheckBus,
68}
69
70impl<const CHUNK: usize, F> BaseAir<F> for PersistentBoundaryAir<CHUNK> {
71    fn width(&self) -> usize {
72        PersistentBoundaryCols::<F, CHUNK>::width()
73    }
74}
75
76impl<const CHUNK: usize, F> BaseAirWithPublicValues<F> for PersistentBoundaryAir<CHUNK> {}
77impl<const CHUNK: usize, F> PartitionedBaseAir<F> for PersistentBoundaryAir<CHUNK> {}
78
79impl<const CHUNK: usize, AB: InteractionBuilder> Air<AB> for PersistentBoundaryAir<CHUNK> {
80    fn eval(&self, builder: &mut AB) {
81        let main = builder.main();
82        let local = main.row_slice(0).expect("window should have two elements");
83        let local: &PersistentBoundaryCols<AB::Var, CHUNK> = (*local).borrow();
84
85        // `direction` should be -1, 0, 1
86        builder.assert_eq(
87            local.expand_direction,
88            local.expand_direction * local.expand_direction * local.expand_direction,
89        );
90
91        // Constrain that an "initial" row has all timestamp zero.
92        // Since `direction` is constrained to be in {-1, 0, 1}, we can select `direction == 1`
93        // with the constraint below.
94        let mut when_initial =
95            builder.when(local.expand_direction * (local.expand_direction + AB::F::ONE));
96        for i in 0..BLOCKS_PER_CHUNK {
97            when_initial.assert_zero(local.timestamps[i]);
98        }
99
100        let mut expand_fields = vec![
101            // direction =  1 => is_final = 0
102            // direction = -1 => is_final = 1
103            local.expand_direction.into(),
104            AB::Expr::ZERO,
105            local.address_space - AB::F::from_u32(ADDR_SPACE_OFFSET),
106            local.leaf_label.into(),
107        ];
108        expand_fields.extend(local.hash.map(Into::into));
109        self.merkle_bus
110            .interact(builder, expand_fields, local.expand_direction.into());
111
112        self.compression_bus.interact(
113            builder,
114            iter::empty()
115                .chain(local.values.map(Into::into))
116                .chain(iter::repeat_n(AB::Expr::ZERO, CHUNK))
117                .chain(local.hash.map(Into::into)),
118            local.expand_direction * local.expand_direction,
119        );
120
121        let chunk_size_f = AB::F::from_usize(CHUNK);
122        for block_idx in 0..BLOCKS_PER_CHUNK {
123            let offset = AB::F::from_usize(block_idx * DEFAULT_BLOCK_SIZE);
124            // Split the 1xCHUNK leaf into DEFAULT_BLOCK_SIZE-sized bus messages.
125            // Each block uses its own timestamp - untouched blocks stay at t=0.
126            self.memory_bus
127                .send(
128                    MemoryAddress::new(
129                        local.address_space,
130                        local.leaf_label * chunk_size_f + offset,
131                    ),
132                    local.values
133                        [block_idx * DEFAULT_BLOCK_SIZE..(block_idx + 1) * DEFAULT_BLOCK_SIZE]
134                        .to_vec(),
135                    local.timestamps[block_idx],
136                )
137                .eval(builder, local.expand_direction);
138        }
139    }
140}
141
142pub struct PersistentBoundaryChip<F, const CHUNK: usize> {
143    pub air: PersistentBoundaryAir<CHUNK>,
144    touched_labels: Option<Vec<FinalTouchedLabel<F, CHUNK>>>,
145    overridden_height: Option<usize>,
146}
147
148#[derive(Debug)]
149pub struct FinalTouchedLabel<F, const CHUNK: usize> {
150    address_space: u32,
151    label: u32,
152    init_values: [F; CHUNK],
153    final_values: [F; CHUNK],
154    init_hash: [F; CHUNK],
155    final_hash: [F; CHUNK],
156    /// Per-block timestamps. Each DEFAULT_BLOCK_SIZE block has its own timestamp.
157    final_timestamps: [u32; BLOCKS_PER_CHUNK],
158}
159
160type BlockInfo<F> = (usize, u32, [F; DEFAULT_BLOCK_SIZE]); // (block_idx, timestamp, values)
161type EnrichedEntry<F> = ((u32, u32), BlockInfo<F>); // (chunk_key, block_info)
162pub(crate) type ChunkedTouchedMemory<F> = Vec<((u32, u32), Vec<BlockInfo<F>>)>;
163
164pub(crate) fn group_touched_memory_by_chunk<F: Copy + Send + Sync>(
165    final_memory: &TimestampedEquipartition<F, DEFAULT_BLOCK_SIZE>,
166) -> ChunkedTouchedMemory<F> {
167    let mut enriched: Vec<EnrichedEntry<F>> = final_memory
168        .par_iter()
169        .map(|&((addr_space, ptr), ts_values)| {
170            let chunk_label = ptr / CHUNK as u32;
171            let block_idx = ((ptr % CHUNK as u32) / DEFAULT_BLOCK_SIZE as u32) as usize;
172            let key = (addr_space, chunk_label);
173            let block_info = (block_idx, ts_values.timestamp, ts_values.values);
174            (key, block_info)
175        })
176        .collect();
177    enriched.sort_unstable_by_key(|(key, _)| *key);
178
179    enriched
180        .chunk_by(|a, b| a.0 == b.0)
181        .map(|group| {
182            let key = group[0].0;
183            let blocks = group.iter().map(|&(_, info)| info).collect();
184            (key, blocks)
185        })
186        .collect()
187}
188
189impl<const CHUNK: usize, F: PrimeField32> PersistentBoundaryChip<F, CHUNK> {
190    pub fn new(
191        memory_bus: MemoryBus,
192        merkle_bus: PermutationCheckBus,
193        compression_bus: PermutationCheckBus,
194    ) -> Self {
195        Self {
196            air: PersistentBoundaryAir {
197                memory_bus,
198                merkle_bus,
199                compression_bus,
200            },
201            touched_labels: None,
202            overridden_height: None,
203        }
204    }
205
206    pub fn set_overridden_height(&mut self, overridden_height: usize) {
207        self.overridden_height = Some(overridden_height);
208    }
209
210    /// Finalize the boundary chip with per-block timestamped memory.
211    ///
212    /// `final_memory` is at DEFAULT_BLOCK_SIZE granularity (4 bytes per entry, single timestamp
213    /// each). This function rechunks into CHUNK-sized (8 bytes) groups with per-block
214    /// timestamps. Untouched blocks within a touched chunk get values from initial_memory and
215    /// timestamp 0.
216    #[instrument(name = "boundary_finalize", level = "debug", skip_all)]
217    pub(crate) fn finalize<H>(
218        &mut self,
219        initial_memory: &MemoryImage,
220        // Touched stuff at DEFAULT_BLOCK_SIZE granularity
221        final_memory: &TimestampedEquipartition<F, DEFAULT_BLOCK_SIZE>,
222        hasher: &H,
223    ) where
224        H: Hasher<CHUNK, F> + Sync + for<'a> SerialReceiver<&'a [F]>,
225    {
226        let final_touched_labels: Vec<_> = group_touched_memory_by_chunk(final_memory)
227            .into_par_iter()
228            .map(|((addr_space, chunk_label), blocks)| {
229                let chunk_ptr = chunk_label * CHUNK as u32;
230                // SAFETY: addr_space from `final_memory` are all in bounds
231                let init_values: [F; CHUNK] = array::from_fn(|i| unsafe {
232                    initial_memory.get_f::<F>(addr_space, chunk_ptr + i as u32)
233                });
234
235                let mut final_values = init_values;
236                let mut timestamps = [0u32; BLOCKS_PER_CHUNK];
237
238                for (block_idx, ts, values) in blocks {
239                    timestamps[block_idx] = ts;
240                    for (i, &val) in values.iter().enumerate() {
241                        final_values[block_idx * DEFAULT_BLOCK_SIZE + i] = val;
242                    }
243                }
244
245                let initial_hash = hasher.hash(&init_values);
246                let final_hash = hasher.hash(&final_values);
247                FinalTouchedLabel {
248                    address_space: addr_space,
249                    label: chunk_label,
250                    init_values,
251                    final_values,
252                    init_hash: initial_hash,
253                    final_hash,
254                    final_timestamps: timestamps,
255                }
256            })
257            .collect();
258        for l in &final_touched_labels {
259            hasher.receive(&l.init_values);
260            hasher.receive(&l.final_values);
261        }
262        self.touched_labels = Some(final_touched_labels);
263    }
264}
265
266impl<const CHUNK: usize, RA, SC> Chip<RA, CpuBackend<SC>> for PersistentBoundaryChip<Val<SC>, CHUNK>
267where
268    SC: StarkProtocolConfig,
269    Val<SC>: PrimeField32,
270{
271    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<CpuBackend<SC>> {
272        let trace = {
273            let touched_labels = self
274                .touched_labels
275                .as_ref()
276                .expect("Cannot generate trace before finalization");
277            let width = PersistentBoundaryCols::<Val<SC>, CHUNK>::width();
278            // Boundary AIR should always present in order to fix the AIR ID of merkle AIR.
279            let mut height = (2 * touched_labels.len()).next_power_of_two();
280            if let Some(mut oh) = self.overridden_height {
281                oh = oh.next_power_of_two();
282                assert!(
283                    oh >= height,
284                    "Overridden height is less than the required height"
285                );
286                height = oh;
287            }
288            let mut rows = Val::<SC>::zero_vec(height * width);
289
290            rows.par_chunks_mut(2 * width)
291                .zip(touched_labels.par_iter())
292                .for_each(|(row, touched_label)| {
293                    let (initial_row, final_row) = row.split_at_mut(width);
294                    *initial_row.borrow_mut() = PersistentBoundaryCols {
295                        expand_direction: Val::<SC>::ONE,
296                        address_space: Val::<SC>::from_u32(touched_label.address_space),
297                        leaf_label: Val::<SC>::from_u32(touched_label.label),
298                        values: touched_label.init_values,
299                        hash: touched_label.init_hash,
300                        timestamps: [Val::<SC>::from_u32(INITIAL_TIMESTAMP); BLOCKS_PER_CHUNK],
301                    };
302
303                    *final_row.borrow_mut() = PersistentBoundaryCols {
304                        expand_direction: Val::<SC>::NEG_ONE,
305                        address_space: Val::<SC>::from_u32(touched_label.address_space),
306                        leaf_label: Val::<SC>::from_u32(touched_label.label),
307                        values: touched_label.final_values,
308                        hash: touched_label.final_hash,
309                        timestamps: touched_label.final_timestamps.map(Val::<SC>::from_u32),
310                    };
311                });
312            RowMajorMatrix::new(rows, width)
313        };
314        AirProvingContext::simple_no_pis(trace)
315    }
316}
openvm_circuit/system/memory/persistent.rs

openvm_circuit/system/memory/
persistent.rs
