openvm_native_circuit/adapters/
loadstore_native_adapter.rs

Help
1use std::{
2    borrow::{Borrow, BorrowMut},
3    marker::PhantomData,
4};
5
6use openvm_circuit::{
7    arch::{
8        instructions::LocalOpcode, AdapterAirContext, AdapterRuntimeContext, ExecutionBridge,
9        ExecutionBus, ExecutionState, Result, VmAdapterAir, VmAdapterChip, VmAdapterInterface,
10    },
11    system::{
12        memory::{
13            offline_checker::{MemoryBridge, MemoryReadAuxCols, MemoryWriteAuxCols},
14            MemoryAddress, MemoryController, OfflineMemory, RecordId,
15        },
16        program::ProgramBus,
17    },
18};
19use openvm_circuit_primitives_derive::AlignedBorrow;
20use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP};
21use openvm_native_compiler::{
22    conversion::AS,
23    NativeLoadStoreOpcode::{self, *},
24};
25use openvm_stark_backend::{
26    interaction::InteractionBuilder,
27    p3_air::BaseAir,
28    p3_field::{Field, FieldAlgebra, PrimeField32},
29};
30use serde::{Deserialize, Serialize};
31
32pub struct NativeLoadStoreInstruction<T> {
33    pub is_valid: T,
34    // Absolute opcode number
35    pub opcode: T,
36    pub is_loadw: T,
37    pub is_storew: T,
38    pub is_hint_storew: T,
39}
40
41pub struct NativeLoadStoreAdapterInterface<T, const NUM_CELLS: usize>(PhantomData<T>);
42
43impl<T, const NUM_CELLS: usize> VmAdapterInterface<T>
44    for NativeLoadStoreAdapterInterface<T, NUM_CELLS>
45{
46    type Reads = (T, [T; NUM_CELLS]);
47    type Writes = [T; NUM_CELLS];
48    type ProcessedInstruction = NativeLoadStoreInstruction<T>;
49}
50
51#[derive(Debug)]
52pub struct NativeLoadStoreAdapterChip<F: Field, const NUM_CELLS: usize> {
53    pub air: NativeLoadStoreAdapterAir<NUM_CELLS>,
54    offset: usize,
55    _marker: PhantomData<F>,
56}
57
58impl<F: PrimeField32, const NUM_CELLS: usize> NativeLoadStoreAdapterChip<F, NUM_CELLS> {
59    pub fn new(
60        execution_bus: ExecutionBus,
61        program_bus: ProgramBus,
62        memory_bridge: MemoryBridge,
63        offset: usize,
64    ) -> Self {
65        Self {
66            air: NativeLoadStoreAdapterAir {
67                memory_bridge,
68                execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
69            },
70            offset,
71            _marker: PhantomData,
72        }
73    }
74}
75
76#[repr(C)]
77#[derive(Clone, Debug, Serialize, Deserialize)]
78#[serde(bound = "F: Field")]
79pub struct NativeLoadStoreReadRecord<F: Field, const NUM_CELLS: usize> {
80    pub pointer_read: RecordId,
81    pub data_read: Option<RecordId>,
82    pub write_as: F,
83    pub write_ptr: F,
84
85    pub a: F,
86    pub b: F,
87    pub c: F,
88    pub d: F,
89    pub e: F,
90}
91
92#[repr(C)]
93#[derive(Clone, Debug, Serialize, Deserialize)]
94#[serde(bound = "F: Field")]
95pub struct NativeLoadStoreWriteRecord<F: Field, const NUM_CELLS: usize> {
96    pub from_state: ExecutionState<F>,
97    pub write_id: RecordId,
98}
99
100#[repr(C)]
101#[derive(Clone, Debug, AlignedBorrow)]
102pub struct NativeLoadStoreAdapterCols<T, const NUM_CELLS: usize> {
103    pub from_state: ExecutionState<T>,
104    pub a: T,
105    pub b: T,
106    pub c: T,
107
108    pub data_write_pointer: T,
109
110    pub pointer_read_aux_cols: MemoryReadAuxCols<T>,
111    pub data_read_aux_cols: MemoryReadAuxCols<T>,
112    pub data_write_aux_cols: MemoryWriteAuxCols<T, NUM_CELLS>,
113}
114
115#[derive(Clone, Copy, Debug, derive_new::new)]
116pub struct NativeLoadStoreAdapterAir<const NUM_CELLS: usize> {
117    pub(super) memory_bridge: MemoryBridge,
118    pub(super) execution_bridge: ExecutionBridge,
119}
120
121impl<F: Field, const NUM_CELLS: usize> BaseAir<F> for NativeLoadStoreAdapterAir<NUM_CELLS> {
122    fn width(&self) -> usize {
123        NativeLoadStoreAdapterCols::<F, NUM_CELLS>::width()
124    }
125}
126
127impl<AB: InteractionBuilder, const NUM_CELLS: usize> VmAdapterAir<AB>
128    for NativeLoadStoreAdapterAir<NUM_CELLS>
129{
130    type Interface = NativeLoadStoreAdapterInterface<AB::Expr, NUM_CELLS>;
131
132    fn eval(
133        &self,
134        builder: &mut AB,
135        local: &[AB::Var],
136        ctx: AdapterAirContext<AB::Expr, Self::Interface>,
137    ) {
138        let cols: &NativeLoadStoreAdapterCols<_, NUM_CELLS> = local.borrow();
139        let timestamp = cols.from_state.timestamp;
140        let mut timestamp_delta = AB::Expr::from_canonical_usize(0);
141
142        let is_valid = ctx.instruction.is_valid;
143        let is_loadw = ctx.instruction.is_loadw;
144        let is_storew = ctx.instruction.is_storew;
145        let is_hint_storew = ctx.instruction.is_hint_storew;
146
147        let native_as = AB::Expr::from_canonical_u32(AS::Native as u32);
148
149        let ptr = ctx.reads.0;
150        // Here we ignore ctx.reads.1 and we use `ctx.writes` as the data for both the write and the
151        // second read (in the case of load/store when it exists).
152        let data = ctx.writes;
153
154        // first pointer read is always [c]_d
155        self.memory_bridge
156            .read(
157                MemoryAddress::new(native_as.clone(), cols.c),
158                [ptr.clone()],
159                timestamp + timestamp_delta.clone(),
160                &cols.pointer_read_aux_cols,
161            )
162            .eval(builder, is_valid.clone());
163        timestamp_delta += is_valid.clone();
164
165        self.memory_bridge
166            .read(
167                MemoryAddress::new(
168                    native_as.clone(),
169                    is_storew.clone() * cols.a + is_loadw.clone() * (ptr.clone() + cols.b),
170                ),
171                data.clone(),
172                timestamp + timestamp_delta.clone(),
173                &cols.data_read_aux_cols,
174            )
175            .eval(builder, is_valid.clone() - is_hint_storew.clone());
176        timestamp_delta += is_valid.clone() - is_hint_storew.clone();
177
178        builder.assert_eq(
179            is_valid.clone() * cols.data_write_pointer,
180            is_loadw.clone() * cols.a
181                + (is_storew.clone() + is_hint_storew.clone()) * (ptr.clone() + cols.b),
182        );
183
184        self.memory_bridge
185            .write(
186                MemoryAddress::new(native_as.clone(), cols.data_write_pointer),
187                data.clone(),
188                timestamp + timestamp_delta.clone(),
189                &cols.data_write_aux_cols,
190            )
191            .eval(builder, is_valid.clone());
192        timestamp_delta += is_valid.clone();
193
194        self.execution_bridge
195            .execute_and_increment_or_set_pc(
196                ctx.instruction.opcode,
197                [
198                    cols.a.into(),
199                    cols.b.into(),
200                    cols.c.into(),
201                    native_as.clone(),
202                    native_as.clone(),
203                ],
204                cols.from_state,
205                timestamp_delta.clone(),
206                (DEFAULT_PC_STEP, ctx.to_pc),
207            )
208            .eval(builder, is_valid.clone());
209    }
210
211    fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
212        let local_cols: &NativeLoadStoreAdapterCols<_, NUM_CELLS> = local.borrow();
213        local_cols.from_state.pc
214    }
215}
216
217impl<F: PrimeField32, const NUM_CELLS: usize> VmAdapterChip<F>
218    for NativeLoadStoreAdapterChip<F, NUM_CELLS>
219{
220    type ReadRecord = NativeLoadStoreReadRecord<F, NUM_CELLS>;
221    type WriteRecord = NativeLoadStoreWriteRecord<F, NUM_CELLS>;
222    type Air = NativeLoadStoreAdapterAir<NUM_CELLS>;
223    type Interface = NativeLoadStoreAdapterInterface<F, NUM_CELLS>;
224
225    fn preprocess(
226        &mut self,
227        memory: &mut MemoryController<F>,
228        instruction: &Instruction<F>,
229    ) -> Result<(
230        <Self::Interface as VmAdapterInterface<F>>::Reads,
231        Self::ReadRecord,
232    )> {
233        let Instruction {
234            opcode,
235            a,
236            b,
237            c,
238            d,
239            e,
240            ..
241        } = *instruction;
242        let local_opcode = NativeLoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.offset));
243
244        let read_as = d;
245        let read_ptr = c;
246        let read_cell = memory.read_cell(read_as, read_ptr);
247
248        let (data_read_as, data_write_as) = {
249            match local_opcode {
250                LOADW => (e, d),
251                STOREW | HINT_STOREW => (d, e),
252            }
253        };
254        let (data_read_ptr, data_write_ptr) = {
255            match local_opcode {
256                LOADW => (read_cell.1 + b, a),
257                STOREW | HINT_STOREW => (a, read_cell.1 + b),
258            }
259        };
260
261        let data_read = match local_opcode {
262            HINT_STOREW => None,
263            LOADW | STOREW => Some(memory.read::<NUM_CELLS>(data_read_as, data_read_ptr)),
264        };
265        let record = NativeLoadStoreReadRecord {
266            pointer_read: read_cell.0,
267            data_read: data_read.map(|x| x.0),
268            write_as: data_write_as,
269            write_ptr: data_write_ptr,
270            a,
271            b,
272            c,
273            d,
274            e,
275        };
276
277        Ok((
278            (read_cell.1, data_read.map_or([F::ZERO; NUM_CELLS], |x| x.1)),
279            record,
280        ))
281    }
282
283    fn postprocess(
284        &mut self,
285        memory: &mut MemoryController<F>,
286        _instruction: &Instruction<F>,
287        from_state: ExecutionState<u32>,
288        output: AdapterRuntimeContext<F, Self::Interface>,
289        read_record: &Self::ReadRecord,
290    ) -> Result<(ExecutionState<u32>, Self::WriteRecord)> {
291        let (write_id, _) =
292            memory.write::<NUM_CELLS>(read_record.write_as, read_record.write_ptr, output.writes);
293        Ok((
294            ExecutionState {
295                pc: output.to_pc.unwrap_or(from_state.pc + DEFAULT_PC_STEP),
296                timestamp: memory.timestamp(),
297            },
298            Self::WriteRecord {
299                from_state: from_state.map(F::from_canonical_u32),
300                write_id,
301            },
302        ))
303    }
304
305    fn generate_trace_row(
306        &self,
307        row_slice: &mut [F],
308        read_record: Self::ReadRecord,
309        write_record: Self::WriteRecord,
310        memory: &OfflineMemory<F>,
311    ) {
312        let aux_cols_factory = memory.aux_cols_factory();
313        let cols: &mut NativeLoadStoreAdapterCols<_, NUM_CELLS> = row_slice.borrow_mut();
314        cols.from_state = write_record.from_state;
315        cols.a = read_record.a;
316        cols.b = read_record.b;
317        cols.c = read_record.c;
318
319        let data_read = read_record.data_read.map(|read| memory.record_by_id(read));
320        if let Some(data_read) = data_read {
321            aux_cols_factory.generate_read_aux(data_read, &mut cols.data_read_aux_cols);
322        }
323
324        let write = memory.record_by_id(write_record.write_id);
325        cols.data_write_pointer = write.pointer;
326
327        aux_cols_factory.generate_read_aux(
328            memory.record_by_id(read_record.pointer_read),
329            &mut cols.pointer_read_aux_cols,
330        );
331        aux_cols_factory.generate_write_aux(write, &mut cols.data_write_aux_cols);
332    }
333
334    fn air(&self) -> &Self::Air {
335        &self.air
336    }
337}
openvm_native_circuit/adapters/loadstore_native_adapter.rs

openvm_native_circuit/adapters/
loadstore_native_adapter.rs
