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#[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}