openvm_native_circuit/loadstore/
core.rs

1use std::{
2    array,
3    borrow::{Borrow, BorrowMut},
4    sync::{Arc, Mutex, OnceLock},
5};
6
7use openvm_circuit::arch::{
8    instructions::LocalOpcode, AdapterAirContext, AdapterRuntimeContext, ExecutionError, Result,
9    Streams, VmAdapterInterface, VmCoreAir, VmCoreChip,
10};
11use openvm_circuit_primitives_derive::AlignedBorrow;
12use openvm_instructions::instruction::Instruction;
13use openvm_native_compiler::NativeLoadStoreOpcode;
14use openvm_stark_backend::{
15    interaction::InteractionBuilder,
16    p3_air::BaseAir,
17    p3_field::{Field, FieldAlgebra, PrimeField32},
18    rap::BaseAirWithPublicValues,
19};
20use serde::{Deserialize, Serialize};
21use serde_big_array::BigArray;
22use strum::IntoEnumIterator;
23
24use super::super::adapters::loadstore_native_adapter::NativeLoadStoreInstruction;
25
26#[repr(C)]
27#[derive(AlignedBorrow)]
28pub struct NativeLoadStoreCoreCols<T, const NUM_CELLS: usize> {
29    pub is_loadw: T,
30    pub is_storew: T,
31    pub is_hint_storew: T,
32
33    pub pointer_read: T,
34    pub data: [T; NUM_CELLS],
35}
36
37#[repr(C)]
38#[derive(Clone, Debug, Serialize, Deserialize)]
39pub struct NativeLoadStoreCoreRecord<F, const NUM_CELLS: usize> {
40    pub opcode: NativeLoadStoreOpcode,
41
42    pub pointer_read: F,
43    #[serde(with = "BigArray")]
44    pub data: [F; NUM_CELLS],
45}
46
47#[derive(Clone, Debug)]
48pub struct NativeLoadStoreCoreAir<const NUM_CELLS: usize> {
49    pub offset: usize,
50}
51
52impl<F: Field, const NUM_CELLS: usize> BaseAir<F> for NativeLoadStoreCoreAir<NUM_CELLS> {
53    fn width(&self) -> usize {
54        NativeLoadStoreCoreCols::<F, NUM_CELLS>::width()
55    }
56}
57
58impl<F: Field, const NUM_CELLS: usize> BaseAirWithPublicValues<F>
59    for NativeLoadStoreCoreAir<NUM_CELLS>
60{
61}
62
63impl<AB, I, const NUM_CELLS: usize> VmCoreAir<AB, I> for NativeLoadStoreCoreAir<NUM_CELLS>
64where
65    AB: InteractionBuilder,
66    I: VmAdapterInterface<AB::Expr>,
67    I::Reads: From<(AB::Expr, [AB::Expr; NUM_CELLS])>,
68    I::Writes: From<[AB::Expr; NUM_CELLS]>,
69    I::ProcessedInstruction: From<NativeLoadStoreInstruction<AB::Expr>>,
70{
71    fn eval(
72        &self,
73        builder: &mut AB,
74        local_core: &[AB::Var],
75        _from_pc: AB::Var,
76    ) -> AdapterAirContext<AB::Expr, I> {
77        let cols: &NativeLoadStoreCoreCols<_, NUM_CELLS> = (*local_core).borrow();
78        let flags = [cols.is_loadw, cols.is_storew, cols.is_hint_storew];
79        let is_valid = flags.iter().fold(AB::Expr::ZERO, |acc, &flag| {
80            builder.assert_bool(flag);
81            acc + flag.into()
82        });
83        builder.assert_bool(is_valid.clone());
84
85        let expected_opcode = VmCoreAir::<AB, I>::expr_to_global_expr(
86            self,
87            flags.iter().zip(NativeLoadStoreOpcode::iter()).fold(
88                AB::Expr::ZERO,
89                |acc, (flag, local_opcode)| {
90                    acc + (*flag).into()
91                        * AB::Expr::from_canonical_usize(local_opcode.local_usize())
92                },
93            ),
94        );
95
96        AdapterAirContext {
97            to_pc: None,
98            reads: (cols.pointer_read.into(), cols.data.map(Into::into)).into(),
99            writes: cols.data.map(Into::into).into(),
100            instruction: NativeLoadStoreInstruction {
101                is_valid,
102                opcode: expected_opcode,
103                is_loadw: cols.is_loadw.into(),
104                is_storew: cols.is_storew.into(),
105                is_hint_storew: cols.is_hint_storew.into(),
106            }
107            .into(),
108        }
109    }
110
111    fn start_offset(&self) -> usize {
112        self.offset
113    }
114}
115
116#[derive(Debug)]
117pub struct NativeLoadStoreCoreChip<F: Field, const NUM_CELLS: usize> {
118    pub air: NativeLoadStoreCoreAir<NUM_CELLS>,
119    pub streams: OnceLock<Arc<Mutex<Streams<F>>>>,
120}
121
122impl<F: Field, const NUM_CELLS: usize> NativeLoadStoreCoreChip<F, NUM_CELLS> {
123    pub fn new(offset: usize) -> Self {
124        Self {
125            air: NativeLoadStoreCoreAir::<NUM_CELLS> { offset },
126            streams: OnceLock::new(),
127        }
128    }
129    pub fn set_streams(&mut self, streams: Arc<Mutex<Streams<F>>>) {
130        self.streams.set(streams).unwrap();
131    }
132}
133
134impl<F: Field, const NUM_CELLS: usize> Default for NativeLoadStoreCoreChip<F, NUM_CELLS> {
135    fn default() -> Self {
136        Self::new(NativeLoadStoreOpcode::CLASS_OFFSET)
137    }
138}
139
140impl<F: PrimeField32, I: VmAdapterInterface<F>, const NUM_CELLS: usize> VmCoreChip<F, I>
141    for NativeLoadStoreCoreChip<F, NUM_CELLS>
142where
143    I::Reads: Into<(F, [F; NUM_CELLS])>,
144    I::Writes: From<[F; NUM_CELLS]>,
145{
146    type Record = NativeLoadStoreCoreRecord<F, NUM_CELLS>;
147    type Air = NativeLoadStoreCoreAir<NUM_CELLS>;
148
149    fn execute_instruction(
150        &self,
151        instruction: &Instruction<F>,
152        from_pc: u32,
153        reads: I::Reads,
154    ) -> Result<(AdapterRuntimeContext<F, I>, Self::Record)> {
155        let Instruction { opcode, .. } = *instruction;
156        let local_opcode =
157            NativeLoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.air.offset));
158        let (pointer_read, data_read) = reads.into();
159
160        let data = if local_opcode == NativeLoadStoreOpcode::HINT_STOREW {
161            let mut streams = self.streams.get().unwrap().lock().unwrap();
162            if streams.hint_stream.len() < NUM_CELLS {
163                return Err(ExecutionError::HintOutOfBounds { pc: from_pc });
164            }
165            array::from_fn(|_| streams.hint_stream.pop_front().unwrap())
166        } else {
167            data_read
168        };
169
170        let output = AdapterRuntimeContext::without_pc(data);
171        let record = NativeLoadStoreCoreRecord {
172            opcode: NativeLoadStoreOpcode::from_usize(opcode.local_opcode_idx(self.air.offset)),
173            pointer_read,
174            data,
175        };
176        Ok((output, record))
177    }
178
179    fn get_opcode_name(&self, opcode: usize) -> String {
180        format!(
181            "{:?}",
182            NativeLoadStoreOpcode::from_usize(opcode - self.air.offset)
183        )
184    }
185
186    fn generate_trace_row(&self, row_slice: &mut [F], record: Self::Record) {
187        let cols: &mut NativeLoadStoreCoreCols<_, NUM_CELLS> = row_slice.borrow_mut();
188        cols.is_loadw = F::from_bool(record.opcode == NativeLoadStoreOpcode::LOADW);
189        cols.is_storew = F::from_bool(record.opcode == NativeLoadStoreOpcode::STOREW);
190        cols.is_hint_storew = F::from_bool(record.opcode == NativeLoadStoreOpcode::HINT_STOREW);
191
192        cols.pointer_read = record.pointer_read;
193        cols.data = record.data;
194    }
195
196    fn air(&self) -> &Self::Air {
197        &self.air
198    }
199}