openvm_circuit/system/memory/
online.rs

1use std::{array::from_fn, fmt::Debug, num::NonZero};
2
3use getset::Getters;
4use itertools::zip_eq;
5use openvm_instructions::exe::SparseMemoryImage;
6use openvm_stark_backend::{
7    p3_field::{Field, PrimeField32},
8    p3_maybe_rayon::prelude::*,
9    p3_util::log2_strict_usize,
10};
11use tracing::instrument;
12
13use crate::{
14    arch::{
15        AddressSpaceHostConfig, AddressSpaceHostLayout, DenseRecordArena, MemoryConfig,
16        RecordArena, MAX_CELL_BYTE_SIZE,
17    },
18    system::{
19        memory::{
20            adapter::records::{AccessLayout, AccessRecordHeader, MERGE_AND_NOT_SPLIT_FLAG},
21            MemoryAddress, TimestampedEquipartition, TimestampedValues, CHUNK,
22        },
23        TouchedMemory,
24    },
25    utils::slice_as_bytes,
26};
27
28mod basic;
29#[cfg(any(unix, windows))]
30mod memmap;
31mod paged_vec;
32
33#[cfg(not(any(unix, windows)))]
34pub use basic::*;
35#[cfg(any(unix, windows))]
36pub use memmap::*;
37pub use paged_vec::PagedVec;
38
39#[cfg(all(any(unix, windows), not(feature = "basic-memory")))]
40pub type MemoryBackend = memmap::MmapMemory;
41#[cfg(any(not(any(unix, windows)), feature = "basic-memory"))]
42pub type MemoryBackend = basic::BasicMemory;
43
44pub const INITIAL_TIMESTAMP: u32 = 0;
45/// Default mmap page size. Change this if using THB.
46pub const PAGE_SIZE: usize = 4096;
47
48// Memory access constraints
49const MAX_BLOCK_SIZE: usize = 32;
50const MIN_ALIGN: usize = 1;
51const MAX_SEGMENTS: usize = MAX_BLOCK_SIZE / MIN_ALIGN;
52
53/// (address_space, pointer)
54pub type Address = (u32, u32);
55
56/// API for any memory implementation that allocates a contiguous region of memory.
57pub trait LinearMemory {
58    /// Create instance of `Self` with `size` bytes.
59    fn new(size: usize) -> Self;
60    /// Allocated size of the memory in bytes.
61    fn size(&self) -> usize;
62    /// Returns the entire memory as a raw byte slice.
63    fn as_slice(&self) -> &[u8];
64    /// Returns the entire memory as a raw byte slice.
65    fn as_mut_slice(&mut self) -> &mut [u8];
66    /// Fill the memory with zeros.
67    fn fill_zero(&mut self) {
68        self.as_mut_slice().fill(0);
69    }
70    /// Read `BLOCK` from `self` at `from` address without moving it.
71    ///
72    /// Panics or segfaults if `from..from + size_of::<BLOCK>()` is out of bounds.
73    ///
74    /// # Safety
75    /// - `BLOCK` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
76    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
77    /// - See [`core::ptr::read`] for similar considerations.
78    /// - Memory at `from` must be properly aligned for `BLOCK`. Use [`Self::read_unaligned`] if
79    ///   alignment is not guaranteed.
80    unsafe fn read<BLOCK: Copy>(&self, from: usize) -> BLOCK;
81    /// Read `BLOCK` from `self` at `from` address without moving it.
82    /// Same as [`Self::read`] except that it does not require alignment.
83    ///
84    /// Panics or segfaults if `from..from + size_of::<BLOCK>()` is out of bounds.
85    ///
86    /// # Safety
87    /// - `BLOCK` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
88    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
89    /// - See [`core::ptr::read`] for similar considerations.
90    unsafe fn read_unaligned<BLOCK: Copy>(&self, from: usize) -> BLOCK;
91    /// Write `BLOCK` to `self` at `start` address without reading the old value. Does not drop
92    /// `values`. Semantically, `values` is moved into the location pointed to by `start`.
93    ///
94    /// Panics or segfaults if `start..start + size_of::<BLOCK>()` is out of bounds.
95    ///
96    /// # Safety
97    /// - See [`core::ptr::write`] for similar considerations.
98    /// - Memory at `start` must be properly aligned for `BLOCK`. Use [`Self::write_unaligned`] if
99    ///   alignment is not guaranteed.
100    unsafe fn write<BLOCK: Copy>(&mut self, start: usize, values: BLOCK);
101    /// Write `BLOCK` to `self` at `start` address without reading the old value. Does not drop
102    /// `values`. Semantically, `values` is moved into the location pointed to by `start`.
103    /// Same as [`Self::write`] but without alignment requirement.
104    ///
105    /// Panics or segfaults if `start..start + size_of::<BLOCK>()` is out of bounds.
106    ///
107    /// # Safety
108    /// - See [`core::ptr::write`] for similar considerations.
109    unsafe fn write_unaligned<BLOCK: Copy>(&mut self, start: usize, values: BLOCK);
110    /// Swaps `values` with memory at `start..start + size_of::<BLOCK>()`.
111    ///
112    /// Panics or segfaults if `start..start + size_of::<BLOCK>()` is out of bounds.
113    ///
114    /// # Safety
115    /// - `BLOCK` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
116    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
117    /// - Memory at `start` must be properly aligned for `BLOCK`.
118    /// - The data in `values` should not overlap with memory in `self`.
119    unsafe fn swap<BLOCK: Copy>(&mut self, start: usize, values: &mut BLOCK);
120    /// Copies `data` into memory at `to` address.
121    ///
122    /// Panics or segfaults if `to..to + size_of_val(data)` is out of bounds.
123    ///
124    /// # Safety
125    /// - `T` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
126    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
127    /// - The underlying memory of `data` should not overlap with `self`.
128    /// - The starting pointer of `self` should be aligned to `T`.
129    /// - The memory pointer at `to` should be aligned to `T`.
130    unsafe fn copy_nonoverlapping<T: Copy>(&mut self, to: usize, data: &[T]);
131    /// Returns a slice `&[T]` for the memory region `start..start + len`.
132    ///
133    /// Panics or segfaults if `start..start + len * size_of::<T>()` is out of bounds.
134    ///
135    /// # Safety
136    /// - `T` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
137    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
138    /// - Memory at `start` must be properly aligned for `T`.
139    unsafe fn get_aligned_slice<T: Copy>(&self, start: usize, len: usize) -> &[T];
140}
141
142/// Map from address space to linear memory.
143/// The underlying memory is typeless, stored as raw bytes, but usage implicitly assumes that each
144/// address space has memory cells of a fixed type (e.g., `u8, F`). We do not use a typemap for
145/// performance reasons, and it is up to the user to enforce types. Needless to say, this is a very
146/// `unsafe` API.
147#[derive(Debug, Clone)]
148#[repr(C)]
149pub struct AddressMap<M: LinearMemory = MemoryBackend> {
150    /// Underlying memory data.
151    pub mem: Vec<M>,
152    /// Host configuration for each address space.
153    pub config: Vec<AddressSpaceHostConfig>,
154}
155
156impl Default for AddressMap {
157    fn default() -> Self {
158        Self::from_mem_config(&MemoryConfig::default())
159    }
160}
161
162impl<M: LinearMemory> AddressMap<M> {
163    pub fn new(config: Vec<AddressSpaceHostConfig>) -> Self {
164        assert_eq!(config[0].num_cells, 0, "Address space 0 must have 0 cells");
165        let mem = config
166            .iter()
167            .map(|config| M::new(config.num_cells.checked_mul(config.layout.size()).unwrap()))
168            .collect();
169        Self { mem, config }
170    }
171
172    pub fn from_mem_config(mem_config: &MemoryConfig) -> Self {
173        Self::new(mem_config.addr_spaces.clone())
174    }
175
176    #[inline(always)]
177    pub fn get_memory(&self) -> &Vec<M> {
178        &self.mem
179    }
180
181    #[inline(always)]
182    pub fn get_memory_mut(&mut self) -> &mut Vec<M> {
183        &mut self.mem
184    }
185
186    /// Fill each address space memory with zeros. Does not change the config.
187    pub fn fill_zero(&mut self) {
188        for mem in &mut self.mem {
189            mem.fill_zero();
190        }
191    }
192
193    /// # Safety
194    /// - Assumes `addr_space` is within the configured memory and not out of bounds
195    pub unsafe fn get_f<F: PrimeField32>(&self, addr_space: u32, ptr: u32) -> F {
196        let layout = &self.config.get_unchecked(addr_space as usize).layout;
197        let start = ptr as usize * layout.size();
198        let bytes = self.get_u8_slice(addr_space, start, layout.size());
199        layout.to_field(bytes)
200    }
201
202    /// # Safety
203    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
204    /// - Assumes `addr_space` is within the configured memory and not out of bounds
205    pub unsafe fn get<T: Copy>(&self, (addr_space, ptr): Address) -> T {
206        debug_assert_eq!(
207            size_of::<T>(),
208            self.config[addr_space as usize].layout.size()
209        );
210        // SAFETY:
211        // - alignment is automatic since we multiply by `size_of::<T>()`
212        self.mem
213            .get_unchecked(addr_space as usize)
214            .read((ptr as usize) * size_of::<T>())
215    }
216
217    /// Panics or segfaults if `ptr..ptr + len` is out of bounds
218    ///
219    /// # Safety
220    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
221    /// - Assumes `addr_space` is within the configured memory and not out of bounds
222    pub unsafe fn get_slice<T: Copy + Debug>(
223        &self,
224        (addr_space, ptr): Address,
225        len: usize,
226    ) -> &[T] {
227        debug_assert_eq!(
228            size_of::<T>(),
229            self.config[addr_space as usize].layout.size()
230        );
231        let start = (ptr as usize) * size_of::<T>();
232        let mem = self.mem.get_unchecked(addr_space as usize);
233        // SAFETY:
234        // - alignment is automatic since we multiply by `size_of::<T>()`
235        mem.get_aligned_slice(start, len)
236    }
237
238    /// Reads the slice at **byte** addresses `start..start + len` from address space `addr_space`
239    /// linear memory. Panics or segfaults if `start..start + len` is out of bounds
240    ///
241    /// # Safety
242    /// - Assumes `addr_space` is within the configured memory and not out of bounds
243    pub unsafe fn get_u8_slice(&self, addr_space: u32, start: usize, len: usize) -> &[u8] {
244        let mem = self.mem.get_unchecked(addr_space as usize);
245        mem.get_aligned_slice(start, len)
246    }
247
248    /// Copies `data` into the memory at `(addr_space, ptr)`.
249    ///
250    /// Panics or segfaults if `ptr + size_of_val(data)` is out of bounds.
251    ///
252    /// # Safety
253    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
254    /// - The linear memory in `addr_space` is aligned to `T`.
255    pub unsafe fn copy_slice_nonoverlapping<T: Copy>(
256        &mut self,
257        (addr_space, ptr): Address,
258        data: &[T],
259    ) {
260        let start = (ptr as usize) * size_of::<T>();
261        // SAFETY:
262        // - Linear memory is aligned to `T` and `start` is multiple of `size_of::<T>()` so
263        //   alignment is satisfied.
264        // - `data` and `self.mem` are non-overlapping
265        self.mem
266            .get_unchecked_mut(addr_space as usize)
267            .copy_nonoverlapping(start, data);
268    }
269
270    // TODO[jpw]: stabilize the boundary memory image format and how to construct
271    /// # Safety
272    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
273    /// - Assumes `addr_space` is within the configured memory and not out of bounds
274    pub fn set_from_sparse(&mut self, sparse_map: &SparseMemoryImage) {
275        for (&(addr_space, index), &data_byte) in sparse_map.iter() {
276            // SAFETY:
277            // - safety assumptions in function doc comments
278            unsafe {
279                self.mem
280                    .get_unchecked_mut(addr_space as usize)
281                    .write_unaligned(index as usize, data_byte);
282            }
283        }
284    }
285}
286
287/// API for guest memory conforming to OpenVM ISA
288// @dev Note we don't make this a trait because phantom executors currently need a concrete type for
289// guest memory
290#[derive(Debug, Clone)]
291#[repr(C)]
292pub struct GuestMemory {
293    pub memory: AddressMap,
294}
295
296impl GuestMemory {
297    pub fn new(addr: AddressMap) -> Self {
298        Self { memory: addr }
299    }
300
301    /// Returns `[pointer:BLOCK_SIZE]_{address_space}`
302    ///
303    /// # Safety
304    /// The type `T` must be stack-allocated `repr(C)` or `repr(transparent)`,
305    /// and it must be the exact type used to represent a single memory cell in
306    /// address space `address_space`. For standard usage,
307    /// `T` is either `u8` or `F` where `F` is the base field of the ZK backend.
308    #[inline(always)]
309    pub unsafe fn read<T, const BLOCK_SIZE: usize>(
310        &self,
311        addr_space: u32,
312        ptr: u32,
313    ) -> [T; BLOCK_SIZE]
314    where
315        T: Copy + Debug,
316    {
317        self.debug_assert_cell_type::<T>(addr_space);
318        // SAFETY:
319        // - `T` should be "plain old data"
320        // - alignment for `[T; BLOCK_SIZE]` is automatic since we multiply by `size_of::<T>()`
321        self.memory
322            .get_memory()
323            .get_unchecked(addr_space as usize)
324            .read((ptr as usize) * size_of::<T>())
325    }
326
327    /// Writes `values` to `[pointer:BLOCK_SIZE]_{address_space}`
328    ///
329    /// # Safety
330    /// See [`GuestMemory::read`].
331    #[inline(always)]
332    pub unsafe fn write<T, const BLOCK_SIZE: usize>(
333        &mut self,
334        addr_space: u32,
335        ptr: u32,
336        values: [T; BLOCK_SIZE],
337    ) where
338        T: Copy + Debug,
339    {
340        self.debug_assert_cell_type::<T>(addr_space);
341        // SAFETY:
342        // - alignment for `[T; BLOCK_SIZE]` is automatic since we multiply by `size_of::<T>()`
343        self.memory
344            .get_memory_mut()
345            .get_unchecked_mut(addr_space as usize)
346            .write((ptr as usize) * size_of::<T>(), values);
347    }
348
349    /// Swaps `values` with `[pointer:BLOCK_SIZE]_{address_space}`.
350    ///
351    /// # Safety
352    /// See [`GuestMemory::read`] and [`LinearMemory::swap`].
353    #[inline(always)]
354    pub unsafe fn swap<T, const BLOCK_SIZE: usize>(
355        &mut self,
356        addr_space: u32,
357        ptr: u32,
358        values: &mut [T; BLOCK_SIZE],
359    ) where
360        T: Copy + Debug,
361    {
362        self.debug_assert_cell_type::<T>(addr_space);
363        // SAFETY:
364        // - alignment for `[T; BLOCK_SIZE]` is automatic since we multiply by `size_of::<T>()`
365        self.memory
366            .get_memory_mut()
367            .get_unchecked_mut(addr_space as usize)
368            .swap((ptr as usize) * size_of::<T>(), values);
369    }
370
371    #[inline(always)]
372    #[allow(clippy::missing_safety_doc)]
373    pub unsafe fn get_slice<T: Copy + Debug>(&self, addr_space: u32, ptr: u32, len: usize) -> &[T] {
374        self.memory.get_slice((addr_space, ptr), len)
375    }
376
377    #[inline(always)]
378    fn debug_assert_cell_type<T>(&self, addr_space: u32) {
379        debug_assert_eq!(
380            size_of::<T>(),
381            self.memory.config[addr_space as usize].layout.size()
382        );
383    }
384}
385
386#[repr(C)]
387#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
388pub struct AccessMetadata {
389    /// Packed timestamp (29 bits) and log2(block_size) (3 bits)
390    pub timestamp_and_log_block_size: u32,
391    /// Offset to block start (in ALIGN units).
392    pub offset_to_start: u8,
393}
394
395impl AccessMetadata {
396    const TIMESTAMP_MASK: u32 = (1 << 29) - 1;
397    const LOG_BLOCK_SIZE_SHIFT: u32 = 29;
398
399    pub fn new(timestamp: u32, block_size: u8, offset_to_start: u8) -> Self {
400        debug_assert!(timestamp < (1 << 29), "Timestamp must be less than 2^29");
401        debug_assert!(
402            block_size == 0 || (block_size.is_power_of_two() && block_size <= MAX_BLOCK_SIZE as u8),
403            "Block size must be 0 or power of 2 and <= {MAX_BLOCK_SIZE}"
404        );
405
406        let encoded_block_size = if block_size == 0 {
407            0
408        } else {
409            // SAFETY: We already asserted that block_size is non-zero in this branch
410            unsafe { NonZero::new_unchecked(block_size) }.ilog2() + 1
411        };
412        let packed = timestamp | (encoded_block_size << Self::LOG_BLOCK_SIZE_SHIFT);
413
414        Self {
415            timestamp_and_log_block_size: packed,
416            offset_to_start,
417        }
418    }
419
420    pub fn timestamp(&self) -> u32 {
421        self.timestamp_and_log_block_size & Self::TIMESTAMP_MASK
422    }
423
424    pub fn block_size(&self) -> u8 {
425        let encoded = self.timestamp_and_log_block_size >> Self::LOG_BLOCK_SIZE_SHIFT;
426        if encoded == 0 {
427            0
428        } else {
429            1 << (encoded - 1)
430        }
431    }
432}
433
434/// Online memory that stores additional information for trace generation purposes.
435/// In particular, keeps track of timestamp.
436#[derive(Getters)]
437pub struct TracingMemory {
438    pub timestamp: u32,
439    /// The initial block size -- this depends on the type of boundary chip.
440    initial_block_size: usize,
441    /// The underlying data memory, with memory cells typed by address space: see [AddressMap].
442    #[getset(get = "pub")]
443    pub data: GuestMemory,
444    /// Maps addr_space to (ptr / min_block_size[addr_space] -> AccessMetadata) for latest access
445    /// metadata. Uses paged storage for memory efficiency. AccessMetadata stores offset_to_start
446    /// (in ALIGN units), block_size, and timestamp (latter two only valid at offset_to_start ==
447    /// 0).
448    pub(super) meta: Vec<PagedVec<AccessMetadata, PAGE_SIZE>>,
449    /// For each `addr_space`, the minimum block size allowed for memory accesses. In other words,
450    /// all memory accesses in `addr_space` must be aligned to this block size.
451    pub min_block_size: Vec<u32>,
452    pub access_adapter_records: DenseRecordArena,
453}
454
455// min_block_size * cell_size never exceeds 8
456const INITIAL_CELL_BUFFER: &[u8] = &[0u8; 8];
457// min_block_size never exceeds 8
458const INITIAL_TIMESTAMP_BUFFER: &[u32] = &[INITIAL_TIMESTAMP; 8];
459
460impl TracingMemory {
461    pub fn new(
462        mem_config: &MemoryConfig,
463        initial_block_size: usize,
464        access_adapter_arena_size_bound: usize,
465    ) -> Self {
466        let image = GuestMemory::new(AddressMap::from_mem_config(mem_config));
467        Self::from_image(image, initial_block_size, access_adapter_arena_size_bound)
468    }
469
470    /// Constructor from pre-existing memory image.
471    pub fn from_image(
472        image: GuestMemory,
473        initial_block_size: usize,
474        access_adapter_arena_size_bound: usize,
475    ) -> Self {
476        let (meta, min_block_size): (Vec<_>, Vec<_>) =
477            zip_eq(image.memory.get_memory(), &image.memory.config)
478                .map(|(mem, addr_sp)| {
479                    let num_cells = mem.size() / addr_sp.layout.size();
480                    let min_block_size = addr_sp.min_block_size;
481                    let total_metadata_len = num_cells.div_ceil(min_block_size);
482                    (PagedVec::new(total_metadata_len), min_block_size as u32)
483                })
484                .unzip();
485        let access_adapter_records =
486            DenseRecordArena::with_byte_capacity(access_adapter_arena_size_bound);
487        Self {
488            data: image,
489            meta,
490            min_block_size,
491            timestamp: INITIAL_TIMESTAMP + 1,
492            initial_block_size,
493            access_adapter_records,
494        }
495    }
496
497    #[inline(always)]
498    fn assert_alignment(&self, block_size: usize, align: usize, addr_space: u32, ptr: u32) {
499        debug_assert!(block_size.is_power_of_two());
500        debug_assert_eq!(block_size % align, 0);
501        debug_assert_ne!(addr_space, 0);
502        debug_assert_eq!(align as u32, self.min_block_size[addr_space as usize]);
503        assert_eq!(
504            ptr % (align as u32),
505            0,
506            "pointer={ptr} not aligned to {align}"
507        );
508    }
509
510    /// Get block metadata by jumping to the start of the block.
511    /// Returns (block_start_pointer, block_metadata).
512    #[inline(always)]
513    fn get_block_metadata<const ALIGN: usize>(
514        &mut self,
515        address_space: usize,
516        pointer: usize,
517    ) -> (u32, AccessMetadata) {
518        let ptr_index = pointer / ALIGN;
519        // SAFETY:
520        // - address_space is validated during instruction decoding and guaranteed to be within
521        //   bounds
522        let meta_page = unsafe { self.meta.get_unchecked_mut(address_space) };
523        let current_meta = meta_page.get(ptr_index);
524
525        let (block_start_index, block_metadata) = if current_meta.offset_to_start == 0 {
526            (ptr_index, current_meta)
527        } else {
528            let offset = current_meta.offset_to_start;
529            let start_idx = ptr_index - offset as usize;
530            let start_meta = meta_page.get(start_idx);
531            (start_idx, start_meta)
532        };
533
534        let block_start_pointer = (block_start_index * ALIGN) as u32;
535
536        (block_start_pointer, block_metadata)
537    }
538
539    #[inline(always)]
540    fn get_timestamp<const ALIGN: usize>(&mut self, address_space: usize, pointer: usize) -> u32 {
541        let ptr_index = pointer / ALIGN;
542        // SAFETY:
543        // - address_space is validated during instruction decoding and guaranteed to be within
544        //   bounds
545        let meta_page = unsafe { self.meta.get_unchecked_mut(address_space) };
546        let current_meta = meta_page.get(ptr_index);
547
548        if current_meta.offset_to_start == 0 {
549            current_meta.timestamp()
550        } else {
551            let offset = current_meta.offset_to_start;
552            let block_start_index = ptr_index - offset as usize;
553            meta_page.get(block_start_index).timestamp()
554        }
555    }
556
557    /// Updates the metadata with the given block.
558    /// Stores timestamp and block_size only at block start, offsets elsewhere.
559    #[inline(always)]
560    fn set_meta_block<const BLOCK_SIZE: usize, const ALIGN: usize>(
561        &mut self,
562        address_space: usize,
563        pointer: usize,
564        timestamp: u32,
565    ) {
566        let ptr = pointer / ALIGN;
567        // SAFETY: address_space is assumed to be valid and within bounds
568        let meta_page = unsafe { self.meta.get_unchecked_mut(address_space) };
569
570        // Store full metadata at the block start
571        meta_page.set(ptr, AccessMetadata::new(timestamp, BLOCK_SIZE as u8, 0));
572
573        // Store offsets for other positions in the block
574        for i in 1..(BLOCK_SIZE / ALIGN) {
575            meta_page.set(ptr + i, AccessMetadata::new(0, 0, i as u8));
576        }
577    }
578
579    pub(crate) fn add_split_record(&mut self, header: AccessRecordHeader) {
580        if header.block_size == header.lowest_block_size {
581            return;
582        }
583        // SAFETY:
584        // - header.address_space is validated during instruction decoding and within bounds
585        // - header.pointer and header.type_size define valid memory bounds within the address space
586        // - The memory access range (header.pointer * header.type_size)..(header.pointer +
587        //   header.block_size) * header.type_size is within the allocated size for the address
588        //   space, preventing out of bounds access
589        let data_slice = unsafe {
590            self.data.memory.get_u8_slice(
591                header.address_space,
592                (header.pointer * header.type_size) as usize,
593                (header.block_size * header.type_size) as usize,
594            )
595        };
596
597        let record_mut = self
598            .access_adapter_records
599            .alloc(AccessLayout::from_record_header(&header));
600        *record_mut.header = header;
601        record_mut.data.copy_from_slice(data_slice);
602        // we don't mind garbage values in prev_*
603    }
604
605    /// `data_slice` is the underlying data of the record in raw host memory format.
606    pub(crate) fn add_merge_record(
607        &mut self,
608        header: AccessRecordHeader,
609        data_slice: &[u8],
610        prev_ts: &[u32],
611    ) {
612        if header.block_size == header.lowest_block_size {
613            return;
614        }
615
616        let record_mut = self
617            .access_adapter_records
618            .alloc(AccessLayout::from_record_header(&header));
619        *record_mut.header = header;
620        record_mut.header.timestamp_and_mask |= MERGE_AND_NOT_SPLIT_FLAG;
621        record_mut.data.copy_from_slice(data_slice);
622        record_mut.timestamps.copy_from_slice(prev_ts);
623    }
624
625    /// Calculate splits and merges needed for a memory access.
626    /// Returns Some((splits, merge)) or None if no operations needed.
627    #[inline(always)]
628    #[allow(clippy::type_complexity)]
629    fn calculate_splits_and_merges<const BLOCK_SIZE: usize, const ALIGN: usize>(
630        &mut self,
631        address_space: usize,
632        pointer: usize,
633    ) -> Option<(Vec<(usize, usize)>, (usize, usize))> {
634        // Skip adapters if this is a repeated access to the same location with same size
635        let (start_ptr, block_meta) = self.get_block_metadata::<ALIGN>(address_space, pointer);
636        if block_meta.block_size() == BLOCK_SIZE as u8 && start_ptr == pointer as u32 {
637            return None;
638        }
639
640        // Split intersecting blocks to align bytes
641        let mut splits_buf = [(0usize, 0usize); MAX_SEGMENTS];
642        let mut splits_count = 0;
643        let mut current_ptr = pointer;
644        let end_ptr = pointer + BLOCK_SIZE;
645
646        while current_ptr < end_ptr {
647            let (start_ptr, block_metadata) =
648                self.get_block_metadata::<ALIGN>(address_space, current_ptr);
649
650            if block_metadata.block_size() == 0 {
651                current_ptr += ALIGN;
652                continue;
653            }
654
655            if block_metadata.block_size() > ALIGN as u8 {
656                // SAFETY: splits_count < MAX_SEGMENTS by construction since we iterate over
657                // at most BLOCK_SIZE/ALIGN segments and BLOCK_SIZE <= MAX_BLOCK_SIZE
658                unsafe {
659                    *splits_buf.get_unchecked_mut(splits_count) =
660                        (start_ptr as usize, block_metadata.block_size() as usize);
661                }
662                splits_count += 1;
663            }
664
665            // Skip to the next segment after this block ends
666            current_ptr = start_ptr as usize + block_metadata.block_size() as usize;
667        }
668
669        let merge = (pointer, BLOCK_SIZE);
670
671        Some((splits_buf[..splits_count].to_vec(), merge))
672    }
673
674    #[inline(always)]
675    fn split_by_meta<T: Copy, const MIN_BLOCK_SIZE: usize>(
676        &mut self,
677        start_ptr: u32,
678        timestamp: u32,
679        block_size: u8,
680        address_space: usize,
681    ) {
682        if block_size == MIN_BLOCK_SIZE as u8 {
683            return;
684        }
685        let begin = start_ptr as usize / MIN_BLOCK_SIZE;
686        // SAFETY:
687        // - address_space is validated during instruction decoding and guaranteed to be within
688        //   bounds
689        let meta_page = unsafe { self.meta.get_unchecked_mut(address_space) };
690
691        for i in 0..(block_size as usize / MIN_BLOCK_SIZE) {
692            // Each split piece becomes its own block start
693            meta_page.set(
694                begin + i,
695                AccessMetadata::new(timestamp, MIN_BLOCK_SIZE as u8, 0),
696            );
697        }
698        self.add_split_record(AccessRecordHeader {
699            timestamp_and_mask: timestamp,
700            address_space: address_space as u32,
701            pointer: start_ptr,
702            block_size: block_size as u32,
703            lowest_block_size: MIN_BLOCK_SIZE as u32,
704            type_size: size_of::<T>() as u32,
705        });
706    }
707
708    /// Returns the timestamp of the previous access to `[pointer:BLOCK_SIZE]_{address_space}`.
709    ///
710    /// Caller must ensure alignment (e.g. via `assert_alignment`) prior to calling this function.
711    #[inline(always)]
712    fn prev_access_time<T: Copy, const BLOCK_SIZE: usize, const ALIGN: usize>(
713        &mut self,
714        address_space: usize,
715        pointer: usize,
716        prev_values: &[T; BLOCK_SIZE],
717    ) -> u32 {
718        debug_assert_eq!(ALIGN, self.data.memory.config[address_space].min_block_size);
719        // SAFETY:
720        // - address_space is validated during instruction decoding and guaranteed to be within
721        //   bounds
722        debug_assert_eq!(
723            unsafe {
724                self.data
725                    .memory
726                    .config
727                    .get_unchecked(address_space)
728                    .layout
729                    .size()
730            },
731            size_of::<T>()
732        );
733        // Calculate what splits and merges are needed for this memory access
734        let result = if let Some((splits, (merge_ptr, merge_size))) =
735            self.calculate_splits_and_merges::<BLOCK_SIZE, ALIGN>(address_space, pointer)
736        {
737            // Process all splits first
738            for (split_ptr, split_size) in splits {
739                let (_, block_metadata) =
740                    self.get_block_metadata::<ALIGN>(address_space, split_ptr);
741                let timestamp = block_metadata.timestamp();
742                self.split_by_meta::<T, ALIGN>(
743                    split_ptr as u32,
744                    timestamp,
745                    split_size as u8,
746                    address_space,
747                );
748            }
749
750            // Process merge
751            let mut prev_ts_buf = [0u32; MAX_SEGMENTS];
752
753            let mut max_timestamp = INITIAL_TIMESTAMP;
754
755            let mut ptr = merge_ptr;
756            let end_ptr = merge_ptr + merge_size;
757            let mut seg_idx = 0;
758            while ptr < end_ptr {
759                let (_, block_metadata) = self.get_block_metadata::<ALIGN>(address_space, ptr);
760
761                let timestamp = if block_metadata.block_size() > 0 {
762                    block_metadata.timestamp()
763                } else {
764                    self.handle_uninitialized_memory::<T, ALIGN>(address_space, ptr);
765                    INITIAL_TIMESTAMP
766                };
767
768                // SAFETY: seg_idx < MAX_SEGMENTS since we iterate at most merge_size/ALIGN times
769                // and merge_size <= BLOCK_SIZE <= MAX_BLOCK_SIZE
770                unsafe {
771                    *prev_ts_buf.get_unchecked_mut(seg_idx) = timestamp;
772                }
773                max_timestamp = max_timestamp.max(timestamp);
774                ptr += ALIGN;
775                seg_idx += 1;
776            }
777
778            // Create the merge record
779            self.add_merge_record(
780                AccessRecordHeader {
781                    timestamp_and_mask: max_timestamp,
782                    address_space: address_space as u32,
783                    pointer: merge_ptr as u32,
784                    block_size: merge_size as u32,
785                    lowest_block_size: ALIGN as u32,
786                    type_size: size_of::<T>() as u32,
787                },
788                // SAFETY: T is plain old data
789                unsafe { slice_as_bytes(prev_values) },
790                &prev_ts_buf[..seg_idx],
791            );
792
793            max_timestamp
794        } else {
795            self.get_timestamp::<ALIGN>(address_space, pointer)
796        };
797
798        // Update the metadata for this access
799        self.set_meta_block::<BLOCK_SIZE, ALIGN>(address_space, pointer, self.timestamp);
800        result
801    }
802
803    /// Handle uninitialized memory by creating appropriate split or merge records.
804    #[inline(always)]
805    fn handle_uninitialized_memory<T: Copy, const ALIGN: usize>(
806        &mut self,
807        address_space: usize,
808        pointer: usize,
809    ) {
810        if self.initial_block_size >= ALIGN {
811            // Split the initial block into chunks
812            let segment_index = pointer / ALIGN;
813            let block_start = segment_index & !(self.initial_block_size / ALIGN - 1);
814            let start_ptr = (block_start * ALIGN) as u32;
815            self.split_by_meta::<T, ALIGN>(
816                start_ptr,
817                INITIAL_TIMESTAMP,
818                self.initial_block_size as u8,
819                address_space,
820            );
821        } else {
822            // Create a merge record for single-byte initialization
823            debug_assert_eq!(self.initial_block_size, 1);
824            self.add_merge_record(
825                AccessRecordHeader {
826                    timestamp_and_mask: INITIAL_TIMESTAMP,
827                    address_space: address_space as u32,
828                    pointer: pointer as u32,
829                    block_size: ALIGN as u32,
830                    lowest_block_size: self.initial_block_size as u32,
831                    type_size: size_of::<T>() as u32,
832                },
833                &INITIAL_CELL_BUFFER[..ALIGN],
834                &INITIAL_TIMESTAMP_BUFFER[..ALIGN],
835            );
836        }
837    }
838
839    /// Atomic read operation which increments the timestamp by 1.
840    /// Returns `(t_prev, [pointer:BLOCK_SIZE]_{address_space})` where `t_prev` is the
841    /// timestamp of the last memory access.
842    ///
843    /// The previous memory access is treated as atomic even if previous accesses were for
844    /// a smaller block size. This is made possible by internal memory access adapters
845    /// that split/merge memory blocks. More specifically, the last memory access corresponding
846    /// to `t_prev` may refer to an atomic access inserted by the memory access adapters.
847    ///
848    /// # Assumptions
849    /// The `BLOCK_SIZE` is a multiple of `ALIGN`, which must equal the minimum block size
850    /// of `address_space`.
851    ///
852    /// # Safety
853    /// The type `T` must be stack-allocated `repr(C)` or `repr(transparent)`,
854    /// plain old data, and it must be the exact type used to represent a single memory cell in
855    /// address space `address_space`. For standard usage,
856    /// `T` is either `u8` or `F` where `F` is the base field of the ZK backend.
857    ///
858    /// In addition:
859    /// - `address_space` must be valid.
860    #[inline(always)]
861    pub unsafe fn read<T, const BLOCK_SIZE: usize, const ALIGN: usize>(
862        &mut self,
863        address_space: u32,
864        pointer: u32,
865    ) -> (u32, [T; BLOCK_SIZE])
866    where
867        T: Copy + Debug,
868    {
869        self.assert_alignment(BLOCK_SIZE, ALIGN, address_space, pointer);
870        let values = self.data.read(address_space, pointer);
871        let t_prev = self.prev_access_time::<T, BLOCK_SIZE, ALIGN>(
872            address_space as usize,
873            pointer as usize,
874            &values,
875        );
876        self.timestamp += 1;
877
878        (t_prev, values)
879    }
880
881    /// Atomic write operation that writes `values` into `[pointer:BLOCK_SIZE]_{address_space}` and
882    /// then increments the timestamp by 1. Returns `(t_prev, values_prev)` which equal the
883    /// timestamp and value `[pointer:BLOCK_SIZE]_{address_space}` of the last memory access.
884    ///
885    /// The previous memory access is treated as atomic even if previous accesses were for
886    /// a smaller block size. This is made possible by internal memory access adapters
887    /// that split/merge memory blocks. More specifically, the last memory access corresponding
888    /// to `t_prev` may refer to an atomic access inserted by the memory access adapters.
889    ///
890    /// # Assumptions
891    /// The `BLOCK_SIZE` is a multiple of `ALIGN`, which must equal the minimum block size
892    /// of `address_space`.
893    ///
894    /// # Safety
895    /// The type `T` must be stack-allocated `repr(C)` or `repr(transparent)`,
896    /// and it must be the exact type used to represent a single memory cell in
897    /// address space `address_space`. For standard usage,
898    /// `T` is either `u8` or `F` where `F` is the base field of the ZK backend.
899    ///
900    /// In addition:
901    /// - `address_space` must be valid.
902    #[inline(always)]
903    pub unsafe fn write<T, const BLOCK_SIZE: usize, const ALIGN: usize>(
904        &mut self,
905        address_space: u32,
906        pointer: u32,
907        values: [T; BLOCK_SIZE],
908    ) -> (u32, [T; BLOCK_SIZE])
909    where
910        T: Copy + Debug,
911    {
912        self.assert_alignment(BLOCK_SIZE, ALIGN, address_space, pointer);
913        let values_prev = self.data.read(address_space, pointer);
914        let t_prev = self.prev_access_time::<T, BLOCK_SIZE, ALIGN>(
915            address_space as usize,
916            pointer as usize,
917            &values_prev,
918        );
919        self.data.write(address_space, pointer, values);
920        self.timestamp += 1;
921
922        (t_prev, values_prev)
923    }
924
925    pub fn increment_timestamp(&mut self) {
926        self.timestamp += 1;
927    }
928
929    pub fn increment_timestamp_by(&mut self, amount: u32) {
930        self.timestamp += amount;
931    }
932
933    pub fn timestamp(&self) -> u32 {
934        self.timestamp
935    }
936
937    /// Finalize the boundary and merkle chips.
938    #[instrument(name = "memory_finalize", skip_all)]
939    pub fn finalize<F: Field>(&mut self, is_persistent: bool) -> TouchedMemory<F> {
940        let touched_blocks = self.touched_blocks();
941
942        match is_persistent {
943            false => TouchedMemory::Volatile(
944                self.touched_blocks_to_equipartition::<F, 1>(touched_blocks),
945            ),
946            true => TouchedMemory::Persistent(
947                self.touched_blocks_to_equipartition::<F, CHUNK>(touched_blocks),
948            ),
949        }
950    }
951
952    /// Returns the list of all touched blocks. The list is sorted by address.
953    fn touched_blocks(&self) -> Vec<(Address, AccessMetadata)> {
954        assert_eq!(self.meta.len(), self.min_block_size.len());
955        self.meta
956            .par_iter()
957            .zip(self.min_block_size.par_iter())
958            .enumerate()
959            .flat_map(|(addr_space, (meta_page, &align))| {
960                meta_page
961                    .par_iter()
962                    .filter_map(move |(idx, metadata)| {
963                        let ptr = idx as u32 * align;
964                        if metadata.offset_to_start == 0 && metadata.block_size() != 0 {
965                            Some(((addr_space as u32, ptr), metadata))
966                        } else {
967                            None
968                        }
969                    })
970                    .collect::<Vec<_>>()
971            })
972            .collect()
973    }
974
975    /// Returns the equipartition of the touched blocks.
976    /// Modifies records and adds new to account for the initial/final segments.
977    fn touched_blocks_to_equipartition<F: Field, const CHUNK: usize>(
978        &mut self,
979        touched_blocks: Vec<((u32, u32), AccessMetadata)>,
980    ) -> TimestampedEquipartition<F, CHUNK> {
981        // [perf] We can `.with_capacity()` if we keep track of the number of segments we initialize
982        let mut final_memory = Vec::new();
983
984        debug_assert!(touched_blocks.is_sorted_by_key(|(addr, _)| addr));
985        self.handle_touched_blocks::<F, CHUNK>(&mut final_memory, touched_blocks);
986
987        debug_assert!(final_memory.is_sorted_by_key(|(key, _)| *key));
988        final_memory
989    }
990
991    fn handle_touched_blocks<F: Field, const CHUNK: usize>(
992        &mut self,
993        final_memory: &mut Vec<((u32, u32), TimestampedValues<F, CHUNK>)>,
994        touched_blocks: Vec<((u32, u32), AccessMetadata)>,
995    ) {
996        let mut current_values = vec![0u8; MAX_CELL_BYTE_SIZE * CHUNK];
997        let mut current_cnt = 0;
998        let mut current_address = MemoryAddress::new(0, 0);
999        let mut current_timestamps = vec![0; CHUNK];
1000        for ((addr_space, ptr), access_metadata) in touched_blocks {
1001            // SAFETY: addr_space of touched blocks are all in bounds
1002            let addr_space_config =
1003                unsafe { *self.data.memory.config.get_unchecked(addr_space as usize) };
1004            let min_block_size = addr_space_config.min_block_size;
1005            let cell_size = addr_space_config.layout.size();
1006            let timestamp = access_metadata.timestamp();
1007            let block_size = access_metadata.block_size();
1008            assert!(
1009                current_cnt == 0
1010                    || (current_address.address_space == addr_space
1011                        && current_address.pointer + current_cnt as u32 == ptr),
1012                "The union of all touched blocks must consist of blocks with sizes divisible by `CHUNK`"
1013            );
1014            debug_assert!(block_size >= min_block_size as u8);
1015            debug_assert!(ptr % min_block_size as u32 == 0);
1016
1017            if current_cnt == 0 {
1018                assert_eq!(
1019                    ptr & (CHUNK as u32 - 1),
1020                    0,
1021                    "The union of all touched blocks must consist of `CHUNK`-aligned blocks"
1022                );
1023                current_address = MemoryAddress::new(addr_space, ptr);
1024            }
1025
1026            if block_size > min_block_size as u8 {
1027                self.add_split_record(AccessRecordHeader {
1028                    timestamp_and_mask: timestamp,
1029                    address_space: addr_space,
1030                    pointer: ptr,
1031                    block_size: block_size as u32,
1032                    lowest_block_size: min_block_size as u32,
1033                    type_size: cell_size as u32,
1034                });
1035            }
1036            if min_block_size > CHUNK {
1037                assert_eq!(current_cnt, 0);
1038                for i in (0..block_size as u32).step_by(min_block_size) {
1039                    self.add_split_record(AccessRecordHeader {
1040                        timestamp_and_mask: timestamp,
1041                        address_space: addr_space,
1042                        pointer: ptr + i,
1043                        block_size: min_block_size as u32,
1044                        lowest_block_size: CHUNK as u32,
1045                        type_size: cell_size as u32,
1046                    });
1047                }
1048                // SAFETY: touched blocks are in bounds
1049                let values = unsafe {
1050                    self.data.memory.get_u8_slice(
1051                        addr_space,
1052                        ptr as usize * cell_size,
1053                        block_size as usize * cell_size,
1054                    )
1055                };
1056                for i in (0..block_size as u32).step_by(CHUNK) {
1057                    final_memory.push((
1058                        (addr_space, ptr + i),
1059                        TimestampedValues {
1060                            timestamp,
1061                            values: from_fn(|j| {
1062                                let byte_idx = (i as usize + j) * cell_size;
1063                                // SAFETY: block_size is multiple of CHUNK and we are reading chunks
1064                                // of cells within bounds
1065                                unsafe {
1066                                    addr_space_config
1067                                        .layout
1068                                        .to_field(&values[byte_idx..byte_idx + cell_size])
1069                                }
1070                            }),
1071                        },
1072                    ));
1073                }
1074            } else {
1075                for i in 0..block_size as u32 {
1076                    // SAFETY: getting cell data
1077                    let cell_data = unsafe {
1078                        self.data.memory.get_u8_slice(
1079                            addr_space,
1080                            (ptr + i) as usize * cell_size,
1081                            cell_size,
1082                        )
1083                    };
1084                    current_values[current_cnt * cell_size..current_cnt * cell_size + cell_size]
1085                        .copy_from_slice(cell_data);
1086                    if current_cnt & (min_block_size - 1) == 0 {
1087                        // SAFETY: current_cnt / min_block_size < CHUNK / min_block_size <= CHUNK
1088                        unsafe {
1089                            *current_timestamps.get_unchecked_mut(current_cnt / min_block_size) =
1090                                timestamp;
1091                        }
1092                    }
1093                    current_cnt += 1;
1094                    if current_cnt == CHUNK {
1095                        let timestamp = *current_timestamps[..CHUNK / min_block_size]
1096                            .iter()
1097                            .max()
1098                            .unwrap();
1099                        self.add_merge_record(
1100                            AccessRecordHeader {
1101                                timestamp_and_mask: timestamp,
1102                                address_space: addr_space,
1103                                pointer: current_address.pointer,
1104                                block_size: CHUNK as u32,
1105                                lowest_block_size: min_block_size as u32,
1106                                type_size: cell_size as u32,
1107                            },
1108                            &current_values[..CHUNK * cell_size],
1109                            &current_timestamps[..CHUNK / min_block_size],
1110                        );
1111                        final_memory.push((
1112                            (current_address.address_space, current_address.pointer),
1113                            TimestampedValues {
1114                                timestamp,
1115                                values: from_fn(|i| unsafe {
1116                                    // SAFETY: cell_size is correct, and alignment is guaranteed
1117                                    addr_space_config.layout.to_field(
1118                                        &current_values[i * cell_size..i * cell_size + cell_size],
1119                                    )
1120                                }),
1121                            },
1122                        ));
1123                        current_address.pointer += current_cnt as u32;
1124                        current_cnt = 0;
1125                    }
1126                }
1127            }
1128        }
1129        assert_eq!(current_cnt, 0, "The union of all touched blocks must consist of blocks with sizes divisible by `CHUNK`");
1130    }
1131
1132    pub fn address_space_alignment(&self) -> Vec<u8> {
1133        self.min_block_size
1134            .iter()
1135            .map(|&x| log2_strict_usize(x as usize) as u8)
1136            .collect()
1137    }
1138}
openvm_circuit/system/memory/online.rs

openvm_circuit/system/memory/
online.rs
