openvm_circuit/system/memory/
online.rs

Help
1use std::fmt::Debug;
2
3use openvm_stark_backend::p3_field::PrimeField32;
4use serde::{Deserialize, Serialize};
5
6use super::paged_vec::{AddressMap, PAGE_SIZE};
7use crate::{
8    arch::MemoryConfig,
9    system::memory::{offline::INITIAL_TIMESTAMP, MemoryImage, RecordId},
10};
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub enum MemoryLogEntry<T> {
14    Read {
15        address_space: u32,
16        pointer: u32,
17        len: usize,
18    },
19    Write {
20        address_space: u32,
21        pointer: u32,
22        data: Vec<T>,
23    },
24    IncrementTimestampBy(u32),
25}
26
27/// A simple data structure to read to/write from memory.
28///
29/// Stores a log of memory accesses to reconstruct aspects of memory state for trace generation.
30#[derive(Debug)]
31pub struct Memory<F> {
32    pub(super) data: AddressMap<F, PAGE_SIZE>,
33    pub(super) log: Vec<MemoryLogEntry<F>>,
34    timestamp: u32,
35}
36
37impl<F: PrimeField32> Memory<F> {
38    pub fn new(mem_config: &MemoryConfig) -> Self {
39        Self {
40            data: AddressMap::from_mem_config(mem_config),
41            timestamp: INITIAL_TIMESTAMP + 1,
42            log: Vec::with_capacity(mem_config.access_capacity),
43        }
44    }
45
46    /// Instantiates a new `Memory` data structure from an image.
47    pub fn from_image(image: MemoryImage<F>, access_capacity: usize) -> Self {
48        Self {
49            data: image,
50            timestamp: INITIAL_TIMESTAMP + 1,
51            log: Vec::with_capacity(access_capacity),
52        }
53    }
54
55    fn last_record_id(&self) -> RecordId {
56        RecordId(self.log.len() - 1)
57    }
58
59    /// Writes an array of values to the memory at the specified address space and start index.
60    ///
61    /// Returns the `RecordId` for the memory record and the previous data.
62    pub fn write<const N: usize>(
63        &mut self,
64        address_space: u32,
65        pointer: u32,
66        values: [F; N],
67    ) -> (RecordId, [F; N]) {
68        assert!(N.is_power_of_two());
69
70        let prev_data = self.data.set_range(&(address_space, pointer), &values);
71
72        self.log.push(MemoryLogEntry::Write {
73            address_space,
74            pointer,
75            data: values.to_vec(),
76        });
77        self.timestamp += 1;
78
79        (self.last_record_id(), prev_data)
80    }
81
82    /// Reads an array of values from the memory at the specified address space and start index.
83    pub fn read<const N: usize>(&mut self, address_space: u32, pointer: u32) -> (RecordId, [F; N]) {
84        assert!(N.is_power_of_two());
85
86        self.log.push(MemoryLogEntry::Read {
87            address_space,
88            pointer,
89            len: N,
90        });
91
92        let values = if address_space == 0 {
93            assert_eq!(N, 1, "cannot batch read from address space 0");
94            [F::from_canonical_u32(pointer); N]
95        } else {
96            self.range_array::<N>(address_space, pointer)
97        };
98        self.timestamp += 1;
99        (self.last_record_id(), values)
100    }
101
102    pub fn increment_timestamp_by(&mut self, amount: u32) {
103        self.timestamp += amount;
104        self.log.push(MemoryLogEntry::IncrementTimestampBy(amount))
105    }
106
107    pub fn timestamp(&self) -> u32 {
108        self.timestamp
109    }
110
111    #[inline(always)]
112    pub fn get(&self, address_space: u32, pointer: u32) -> F {
113        *self.data.get(&(address_space, pointer)).unwrap_or(&F::ZERO)
114    }
115
116    #[inline(always)]
117    fn range_array<const N: usize>(&self, address_space: u32, pointer: u32) -> [F; N] {
118        self.data.get_range(&(address_space, pointer))
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use openvm_stark_backend::p3_field::FieldAlgebra;
125    use openvm_stark_sdk::p3_baby_bear::BabyBear;
126
127    use super::Memory;
128    use crate::arch::MemoryConfig;
129
130    macro_rules! bba {
131        [$($x:expr),*] => {
132            [$(BabyBear::from_canonical_u32($x)),*]
133        }
134    }
135
136    #[test]
137    fn test_write_read() {
138        let mut memory = Memory::new(&MemoryConfig::default());
139        let address_space = 1;
140
141        memory.write(address_space, 0, bba![1, 2, 3, 4]);
142
143        let (_, data) = memory.read::<2>(address_space, 0);
144        assert_eq!(data, bba![1, 2]);
145
146        memory.write(address_space, 2, bba![100]);
147
148        let (_, data) = memory.read::<4>(address_space, 0);
149        assert_eq!(data, bba![1, 2, 100, 4]);
150    }
151}
openvm_circuit/system/memory/online.rs

openvm_circuit/system/memory/
online.rs
