openvm_circuit/arch/testing/memory/
mod.rs

1use std::collections::HashMap;
2
3use air::{MemoryDummyAir, MemoryDummyChip};
4use rand::Rng;
5
6use crate::{
7    arch::VmField,
8    system::memory::{online::TracingMemory, MemoryController},
9};
10
11pub mod air;
12
13#[cfg(feature = "cuda")]
14mod cuda;
15#[cfg(feature = "cuda")]
16pub use cuda::*;
17
18/// A dummy testing chip that will add unconstrained messages into the [MemoryBus].
19/// Stores a log of raw messages to send/receive to the [MemoryBus].
20///
21/// It will create a [air::MemoryDummyAir] to add messages to MemoryBus.
22pub struct MemoryTester<F: VmField> {
23    /// Map from `block_size` to [MemoryDummyChip] of that block size
24    pub chip_for_block: HashMap<usize, MemoryDummyChip<F>>,
25    pub memory: TracingMemory,
26    pub(super) controller: MemoryController<F>,
27}
28
29impl<F: VmField> MemoryTester<F> {
30    pub fn new(controller: MemoryController<F>, memory: TracingMemory) -> Self {
31        let bus = controller.memory_bus;
32        let mut chip_for_block = HashMap::new();
33        for log_block_size in 0..6 {
34            let block_size = 1 << log_block_size;
35            let chip = MemoryDummyChip::new(MemoryDummyAir::new(bus, block_size));
36            chip_for_block.insert(block_size, chip);
37        }
38        Self {
39            chip_for_block,
40            memory,
41            controller,
42        }
43    }
44
45    pub fn read<const N: usize>(&mut self, addr_space: usize, ptr: usize) -> [F; N] {
46        let memory = &mut self.memory;
47        let t = memory.timestamp();
48        // TODO: this could be improved if we added a TracingMemory::get_f function
49        let (t_prev, data) = if addr_space <= 3 {
50            let (t_prev, data) = unsafe { memory.read::<u8, N, 4>(addr_space as u32, ptr as u32) };
51            (t_prev, data.map(F::from_u8))
52        } else {
53            unsafe { memory.read::<F, N, 1>(addr_space as u32, ptr as u32) }
54        };
55        self.chip_for_block.get_mut(&N).unwrap().receive(
56            addr_space as u32,
57            ptr as u32,
58            &data,
59            t_prev,
60        );
61        self.chip_for_block
62            .get_mut(&N)
63            .unwrap()
64            .send(addr_space as u32, ptr as u32, &data, t);
65
66        data
67    }
68
69    pub fn write<const N: usize>(&mut self, addr_space: usize, ptr: usize, data: [F; N]) {
70        let memory = &mut self.memory;
71        let t = memory.timestamp();
72        // TODO: this could be improved if we added a TracingMemory::write_f function
73        let (t_prev, data_prev) = if addr_space <= 3 {
74            let (t_prev, data_prev) = unsafe {
75                memory.write::<u8, N, 4>(
76                    addr_space as u32,
77                    ptr as u32,
78                    data.map(|x| x.as_canonical_u32() as u8),
79                )
80            };
81            (t_prev, data_prev.map(F::from_u8))
82        } else {
83            unsafe { memory.write::<F, N, 1>(addr_space as u32, ptr as u32, data) }
84        };
85        self.chip_for_block.get_mut(&N).unwrap().receive(
86            addr_space as u32,
87            ptr as u32,
88            &data_prev,
89            t_prev,
90        );
91        self.chip_for_block
92            .get_mut(&N)
93            .unwrap()
94            .send(addr_space as u32, ptr as u32, &data, t);
95    }
96}
97
98pub fn gen_pointer<R>(rng: &mut R, len: usize) -> usize
99where
100    R: Rng + ?Sized,
101{
102    const MAX_MEMORY: usize = 1 << 29;
103    rng.random_range(0..MAX_MEMORY - len) / len * len
104}
openvm_circuit/arch/testing/memory/mod.rs

openvm_circuit/arch/testing/memory/
mod.rs
