openvm_circuit/arch/testing/memory/
air.rs

1use std::{mem::size_of, sync::Arc};
2
3use openvm_stark_backend::{
4    config::{StarkGenericConfig, Val},
5    interaction::InteractionBuilder,
6    p3_air::{Air, BaseAir},
7    p3_field::{FieldAlgebra, PrimeField32},
8    p3_matrix::{dense::RowMajorMatrix, Matrix},
9    prover::{cpu::CpuBackend, types::AirProvingContext},
10    rap::{BaseAirWithPublicValues, PartitionedBaseAir},
11    Chip, ChipUsageGetter,
12};
13
14use crate::system::memory::{offline_checker::MemoryBus, MemoryAddress};
15
16#[repr(C)]
17#[derive(Clone, Copy)]
18pub struct DummyMemoryInteractionColsRef<'a, T> {
19    pub address: MemoryAddress<&'a T, &'a T>,
20    pub data: &'a [T],
21    pub timestamp: &'a T,
22    /// The send frequency. Send corresponds to write. To read, set to negative.
23    pub count: &'a T,
24}
25
26#[repr(C)]
27pub struct DummyMemoryInteractionColsMut<'a, T> {
28    pub address: MemoryAddress<&'a mut T, &'a mut T>,
29    pub data: &'a mut [T],
30    pub timestamp: &'a mut T,
31    /// The send frequency. Send corresponds to write. To read, set to negative.
32    pub count: &'a mut T,
33}
34
35impl<'a, T> DummyMemoryInteractionColsRef<'a, T> {
36    pub fn from_slice(slice: &'a [T]) -> Self {
37        let (address, slice) = slice.split_at(size_of::<MemoryAddress<u8, u8>>());
38        let (count, slice) = slice.split_last().unwrap();
39        let (timestamp, data) = slice.split_last().unwrap();
40        Self {
41            address: MemoryAddress::new(&address[0], &address[1]),
42            data,
43            timestamp,
44            count,
45        }
46    }
47}
48
49impl<'a, T> DummyMemoryInteractionColsMut<'a, T> {
50    pub fn from_mut_slice(slice: &'a mut [T]) -> Self {
51        let (addr_space, slice) = slice.split_first_mut().unwrap();
52        let (ptr, slice) = slice.split_first_mut().unwrap();
53        let (count, slice) = slice.split_last_mut().unwrap();
54        let (timestamp, data) = slice.split_last_mut().unwrap();
55        Self {
56            address: MemoryAddress::new(addr_space, ptr),
57            data,
58            timestamp,
59            count,
60        }
61    }
62}
63
64#[derive(Clone, Copy, Debug, derive_new::new)]
65pub struct MemoryDummyAir {
66    pub bus: MemoryBus,
67    pub block_size: usize,
68}
69
70impl<F> BaseAirWithPublicValues<F> for MemoryDummyAir {}
71impl<F> PartitionedBaseAir<F> for MemoryDummyAir {}
72impl<F> BaseAir<F> for MemoryDummyAir {
73    fn width(&self) -> usize {
74        self.block_size + 4
75    }
76}
77
78impl<AB: InteractionBuilder> Air<AB> for MemoryDummyAir {
79    fn eval(&self, builder: &mut AB) {
80        let main = builder.main();
81        let local = main.row_slice(0);
82        let local = DummyMemoryInteractionColsRef::from_slice(&local);
83
84        self.bus
85            .send(
86                MemoryAddress::new(*local.address.address_space, *local.address.pointer),
87                local.data.to_vec(),
88                *local.timestamp,
89            )
90            .eval(builder, *local.count);
91    }
92}
93
94#[derive(Clone)]
95pub struct MemoryDummyChip<F> {
96    pub air: MemoryDummyAir,
97    pub trace: Vec<F>,
98}
99
100impl<F> MemoryDummyChip<F> {
101    pub fn new(air: MemoryDummyAir) -> Self {
102        Self {
103            air,
104            trace: Vec::new(),
105        }
106    }
107}
108
109impl<F: PrimeField32> MemoryDummyChip<F> {
110    pub fn send(&mut self, addr_space: u32, ptr: u32, data: &[F], timestamp: u32) {
111        self.push(addr_space, ptr, data, timestamp, F::ONE);
112    }
113
114    pub fn receive(&mut self, addr_space: u32, ptr: u32, data: &[F], timestamp: u32) {
115        self.push(addr_space, ptr, data, timestamp, F::NEG_ONE);
116    }
117
118    pub fn push(&mut self, addr_space: u32, ptr: u32, data: &[F], timestamp: u32, count: F) {
119        assert_eq!(data.len(), self.air.block_size);
120        self.trace.push(F::from_canonical_u32(addr_space));
121        self.trace.push(F::from_canonical_u32(ptr));
122        self.trace.extend_from_slice(data);
123        self.trace.push(F::from_canonical_u32(timestamp));
124        self.trace.push(count);
125    }
126}
127
128impl<SC: StarkGenericConfig, RA> Chip<RA, CpuBackend<SC>> for MemoryDummyChip<Val<SC>>
129where
130    Val<SC>: PrimeField32,
131{
132    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<CpuBackend<SC>> {
133        let height = self.current_trace_height().next_power_of_two();
134        let width = self.trace_width();
135        let mut trace = self.trace.clone();
136        trace.resize(height * width, Val::<SC>::ZERO);
137
138        let trace = Arc::new(RowMajorMatrix::new(trace, width));
139        AirProvingContext::simple_no_pis(trace)
140    }
141}
142
143impl<F: PrimeField32> ChipUsageGetter for MemoryDummyChip<F> {
144    fn air_name(&self) -> String {
145        format!("MemoryDummyAir<{}>", self.air.block_size)
146    }
147    fn current_trace_height(&self) -> usize {
148        self.trace.len() / self.trace_width()
149    }
150    fn trace_width(&self) -> usize {
151        BaseAir::<F>::width(&self.air)
152    }
153}
openvm_circuit/arch/testing/memory/air.rs

openvm_circuit/arch/testing/memory/
air.rs
