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

openvm_circuit/system/memory/
online.rs
