openvm_circuit/system/native_adapter/
mod.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    marker::PhantomData,
4};
5
6use openvm_circuit::{
7    arch::{
8        AdapterAirContext, AdapterRuntimeContext, BasicAdapterInterface, ExecutionBridge,
9        ExecutionBus, ExecutionState, MinimalInstruction, Result, VmAdapterAir, VmAdapterChip,
10        VmAdapterInterface,
11    },
12    system::{
13        memory::{
14            offline_checker::{MemoryBridge, MemoryReadOrImmediateAuxCols, MemoryWriteAuxCols},
15            MemoryAddress, MemoryController,
16        },
17        program::ProgramBus,
18    },
19};
20use openvm_circuit_primitives_derive::AlignedBorrow;
21use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP};
22use openvm_stark_backend::{
23    interaction::InteractionBuilder,
24    p3_air::BaseAir,
25    p3_field::{Field, FieldAlgebra, PrimeField32},
26};
27use serde::{Deserialize, Serialize};
28use serde_big_array::BigArray;
29
30use crate::system::memory::{OfflineMemory, RecordId};
31
32/// R reads(R<=2), W writes(W<=1).
33/// Operands: b for the first read, c for the second read, a for the first write.
34/// If an operand is not used, its address space and pointer should be all 0.
35#[derive(Debug)]
36pub struct NativeAdapterChip<F, const R: usize, const W: usize> {
37    pub air: NativeAdapterAir<R, W>,
38    _phantom: PhantomData<F>,
39}
40
41impl<F: PrimeField32, const R: usize, const W: usize> NativeAdapterChip<F, R, W> {
42    pub fn new(
43        execution_bus: ExecutionBus,
44        program_bus: ProgramBus,
45        memory_bridge: MemoryBridge,
46    ) -> Self {
47        Self {
48            air: NativeAdapterAir {
49                execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
50                memory_bridge,
51            },
52            _phantom: PhantomData,
53        }
54    }
55}
56
57#[repr(C)]
58#[derive(Debug, Serialize, Deserialize)]
59pub struct NativeReadRecord<F: Field, const R: usize> {
60    #[serde(with = "BigArray")]
61    pub reads: [(RecordId, [F; 1]); R],
62}
63
64impl<F: Field, const R: usize> NativeReadRecord<F, R> {
65    pub fn b(&self) -> &[F; 1] {
66        &self.reads[0].1
67    }
68
69    pub fn c(&self) -> &[F; 1] {
70        &self.reads[1].1
71    }
72}
73
74#[repr(C)]
75#[derive(Debug, Serialize, Deserialize)]
76#[serde(bound = "F: Field")]
77pub struct NativeWriteRecord<F: Field, const W: usize> {
78    pub from_state: ExecutionState<u32>,
79    #[serde(with = "BigArray")]
80    pub writes: [(RecordId, [F; 1]); W],
81}
82
83impl<F: Field, const W: usize> NativeWriteRecord<F, W> {
84    pub fn a(&self) -> &[F; 1] {
85        &self.writes[0].1
86    }
87}
88
89#[repr(C)]
90#[derive(AlignedBorrow)]
91pub struct NativeAdapterReadCols<T> {
92    pub address: MemoryAddress<T, T>,
93    pub read_aux: MemoryReadOrImmediateAuxCols<T>,
94}
95
96#[repr(C)]
97#[derive(AlignedBorrow)]
98pub struct NativeAdapterWriteCols<T> {
99    pub address: MemoryAddress<T, T>,
100    pub write_aux: MemoryWriteAuxCols<T, 1>,
101}
102
103#[repr(C)]
104#[derive(AlignedBorrow)]
105pub struct NativeAdapterCols<T, const R: usize, const W: usize> {
106    pub from_state: ExecutionState<T>,
107    pub reads_aux: [NativeAdapterReadCols<T>; R],
108    pub writes_aux: [NativeAdapterWriteCols<T>; W],
109}
110
111#[derive(Clone, Copy, Debug, derive_new::new)]
112pub struct NativeAdapterAir<const R: usize, const W: usize> {
113    pub(super) execution_bridge: ExecutionBridge,
114    pub(super) memory_bridge: MemoryBridge,
115}
116
117impl<F: Field, const R: usize, const W: usize> BaseAir<F> for NativeAdapterAir<R, W> {
118    fn width(&self) -> usize {
119        NativeAdapterCols::<F, R, W>::width()
120    }
121}
122
123impl<AB: InteractionBuilder, const R: usize, const W: usize> VmAdapterAir<AB>
124    for NativeAdapterAir<R, W>
125{
126    type Interface = BasicAdapterInterface<AB::Expr, MinimalInstruction<AB::Expr>, R, W, 1, 1>;
127
128    fn eval(
129        &self,
130        builder: &mut AB,
131        local: &[AB::Var],
132        ctx: AdapterAirContext<AB::Expr, Self::Interface>,
133    ) {
134        let cols: &NativeAdapterCols<_, R, W> = local.borrow();
135        let timestamp = cols.from_state.timestamp;
136        let mut timestamp_delta = 0usize;
137        let mut timestamp_pp = || {
138            timestamp_delta += 1;
139            timestamp + AB::F::from_canonical_usize(timestamp_delta - 1)
140        };
141
142        for (i, r_cols) in cols.reads_aux.iter().enumerate() {
143            self.memory_bridge
144                .read_or_immediate(
145                    r_cols.address,
146                    ctx.reads[i][0].clone(),
147                    timestamp_pp(),
148                    &r_cols.read_aux,
149                )
150                .eval(builder, ctx.instruction.is_valid.clone());
151        }
152        for (i, w_cols) in cols.writes_aux.iter().enumerate() {
153            self.memory_bridge
154                .write(
155                    w_cols.address,
156                    ctx.writes[i].clone(),
157                    timestamp_pp(),
158                    &w_cols.write_aux,
159                )
160                .eval(builder, ctx.instruction.is_valid.clone());
161        }
162
163        let zero_address =
164            || MemoryAddress::new(AB::Expr::from(AB::F::ZERO), AB::Expr::from(AB::F::ZERO));
165        let f = |var_addr: MemoryAddress<AB::Var, AB::Var>| -> MemoryAddress<AB::Expr, AB::Expr> {
166            MemoryAddress::new(var_addr.address_space.into(), var_addr.pointer.into())
167        };
168
169        let addr_a = if W >= 1 {
170            f(cols.writes_aux[0].address)
171        } else {
172            zero_address()
173        };
174        let addr_b = if R >= 1 {
175            f(cols.reads_aux[0].address)
176        } else {
177            zero_address()
178        };
179        let addr_c = if R >= 2 {
180            f(cols.reads_aux[1].address)
181        } else {
182            zero_address()
183        };
184        self.execution_bridge
185            .execute_and_increment_or_set_pc(
186                ctx.instruction.opcode,
187                [
188                    addr_a.pointer,
189                    addr_b.pointer,
190                    addr_c.pointer,
191                    addr_a.address_space,
192                    addr_b.address_space,
193                    addr_c.address_space,
194                ],
195                cols.from_state,
196                AB::F::from_canonical_usize(timestamp_delta),
197                (DEFAULT_PC_STEP, ctx.to_pc),
198            )
199            .eval(builder, ctx.instruction.is_valid);
200    }
201
202    fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
203        let cols: &NativeAdapterCols<_, R, W> = local.borrow();
204        cols.from_state.pc
205    }
206}
207
208impl<F: PrimeField32, const R: usize, const W: usize> VmAdapterChip<F>
209    for NativeAdapterChip<F, R, W>
210{
211    type ReadRecord = NativeReadRecord<F, R>;
212    type WriteRecord = NativeWriteRecord<F, W>;
213    type Air = NativeAdapterAir<R, W>;
214    type Interface = BasicAdapterInterface<F, MinimalInstruction<F>, R, W, 1, 1>;
215
216    fn preprocess(
217        &mut self,
218        memory: &mut MemoryController<F>,
219        instruction: &Instruction<F>,
220    ) -> Result<(
221        <Self::Interface as VmAdapterInterface<F>>::Reads,
222        Self::ReadRecord,
223    )> {
224        assert!(R <= 2);
225        let Instruction { b, c, e, f, .. } = *instruction;
226
227        let mut reads = Vec::with_capacity(R);
228        if R >= 1 {
229            reads.push(memory.read::<1>(e, b));
230        }
231        if R >= 2 {
232            reads.push(memory.read::<1>(f, c));
233        }
234        let i_reads: [_; R] = std::array::from_fn(|i| reads[i].1);
235
236        Ok((
237            i_reads,
238            Self::ReadRecord {
239                reads: reads.try_into().unwrap(),
240            },
241        ))
242    }
243
244    fn postprocess(
245        &mut self,
246        memory: &mut MemoryController<F>,
247        instruction: &Instruction<F>,
248        from_state: ExecutionState<u32>,
249        output: AdapterRuntimeContext<F, Self::Interface>,
250        _read_record: &Self::ReadRecord,
251    ) -> Result<(ExecutionState<u32>, Self::WriteRecord)> {
252        assert!(W <= 1);
253        let Instruction { a, d, .. } = *instruction;
254        let mut writes = Vec::with_capacity(W);
255        if W >= 1 {
256            let (record_id, _) = memory.write(d, a, output.writes[0]);
257            writes.push((record_id, output.writes[0]));
258        }
259
260        Ok((
261            ExecutionState {
262                pc: output.to_pc.unwrap_or(from_state.pc + DEFAULT_PC_STEP),
263                timestamp: memory.timestamp(),
264            },
265            Self::WriteRecord {
266                from_state,
267                writes: writes.try_into().unwrap(),
268            },
269        ))
270    }
271
272    fn generate_trace_row(
273        &self,
274        row_slice: &mut [F],
275        read_record: Self::ReadRecord,
276        write_record: Self::WriteRecord,
277        memory: &OfflineMemory<F>,
278    ) {
279        let row_slice: &mut NativeAdapterCols<_, R, W> = row_slice.borrow_mut();
280        let aux_cols_factory = memory.aux_cols_factory();
281
282        row_slice.from_state = write_record.from_state.map(F::from_canonical_u32);
283
284        for (i, read) in read_record.reads.iter().enumerate() {
285            let (id, _) = read;
286            let record = memory.record_by_id(*id);
287            aux_cols_factory
288                .generate_read_or_immediate_aux(record, &mut row_slice.reads_aux[i].read_aux);
289            row_slice.reads_aux[i].address =
290                MemoryAddress::new(record.address_space, record.pointer);
291        }
292
293        for (i, write) in write_record.writes.iter().enumerate() {
294            let (id, _) = write;
295            let record = memory.record_by_id(*id);
296            aux_cols_factory.generate_write_aux(record, &mut row_slice.writes_aux[i].write_aux);
297            row_slice.writes_aux[i].address =
298                MemoryAddress::new(record.address_space, record.pointer);
299        }
300    }
301
302    fn air(&self) -> &Self::Air {
303        &self.air
304    }
305}