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, MemoryReadAuxCols, MemoryWriteAuxCols},
15 MemoryAddress, MemoryController, OfflineMemory, RecordId,
16 },
17 program::ProgramBus,
18 },
19};
20use openvm_circuit_primitives_derive::AlignedBorrow;
21use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP};
22use openvm_native_compiler::conversion::AS;
23use openvm_stark_backend::{
24 interaction::InteractionBuilder,
25 p3_air::BaseAir,
26 p3_field::{Field, FieldAlgebra, PrimeField32},
27};
28use serde::{Deserialize, Serialize};
29
30#[allow(dead_code)]
31#[derive(Debug)]
32pub struct NativeVectorizedAdapterChip<F: Field, const N: usize> {
33 pub air: NativeVectorizedAdapterAir<N>,
34 _marker: PhantomData<F>,
35}
36
37impl<F: PrimeField32, const N: usize> NativeVectorizedAdapterChip<F, N> {
38 pub fn new(
39 execution_bus: ExecutionBus,
40 program_bus: ProgramBus,
41 memory_bridge: MemoryBridge,
42 ) -> Self {
43 Self {
44 air: NativeVectorizedAdapterAir {
45 execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
46 memory_bridge,
47 },
48 _marker: PhantomData,
49 }
50 }
51}
52
53#[repr(C)]
54#[derive(Debug, Serialize, Deserialize)]
55pub struct NativeVectorizedReadRecord<const N: usize> {
56 pub b: RecordId,
57 pub c: RecordId,
58}
59
60#[repr(C)]
61#[derive(Debug, Serialize, Deserialize)]
62pub struct NativeVectorizedWriteRecord<const N: usize> {
63 pub from_state: ExecutionState<u32>,
64 pub a: RecordId,
65}
66
67#[repr(C)]
68#[derive(AlignedBorrow)]
69pub struct NativeVectorizedAdapterCols<T, const N: usize> {
70 pub from_state: ExecutionState<T>,
71 pub a_pointer: T,
72 pub b_pointer: T,
73 pub c_pointer: T,
74 pub reads_aux: [MemoryReadAuxCols<T>; 2],
75 pub writes_aux: [MemoryWriteAuxCols<T, N>; 1],
76}
77
78#[derive(Clone, Copy, Debug, derive_new::new)]
79pub struct NativeVectorizedAdapterAir<const N: usize> {
80 pub(super) execution_bridge: ExecutionBridge,
81 pub(super) memory_bridge: MemoryBridge,
82}
83
84impl<F: Field, const N: usize> BaseAir<F> for NativeVectorizedAdapterAir<N> {
85 fn width(&self) -> usize {
86 NativeVectorizedAdapterCols::<F, N>::width()
87 }
88}
89
90impl<AB: InteractionBuilder, const N: usize> VmAdapterAir<AB> for NativeVectorizedAdapterAir<N> {
91 type Interface = BasicAdapterInterface<AB::Expr, MinimalInstruction<AB::Expr>, 2, 1, N, N>;
92
93 fn eval(
94 &self,
95 builder: &mut AB,
96 local: &[AB::Var],
97 ctx: AdapterAirContext<AB::Expr, Self::Interface>,
98 ) {
99 let cols: &NativeVectorizedAdapterCols<_, N> = local.borrow();
100 let timestamp = cols.from_state.timestamp;
101 let mut timestamp_delta = 0usize;
102 let mut timestamp_pp = || {
103 timestamp_delta += 1;
104 timestamp + AB::F::from_canonical_usize(timestamp_delta - 1)
105 };
106
107 let native_as = AB::Expr::from_canonical_u32(AS::Native as u32);
108
109 self.memory_bridge
110 .read(
111 MemoryAddress::new(native_as.clone(), cols.b_pointer),
112 ctx.reads[0].clone(),
113 timestamp_pp(),
114 &cols.reads_aux[0],
115 )
116 .eval(builder, ctx.instruction.is_valid.clone());
117
118 self.memory_bridge
119 .read(
120 MemoryAddress::new(native_as.clone(), cols.c_pointer),
121 ctx.reads[1].clone(),
122 timestamp_pp(),
123 &cols.reads_aux[1],
124 )
125 .eval(builder, ctx.instruction.is_valid.clone());
126
127 self.memory_bridge
128 .write(
129 MemoryAddress::new(native_as.clone(), cols.a_pointer),
130 ctx.writes[0].clone(),
131 timestamp_pp(),
132 &cols.writes_aux[0],
133 )
134 .eval(builder, ctx.instruction.is_valid.clone());
135
136 self.execution_bridge
137 .execute_and_increment_or_set_pc(
138 ctx.instruction.opcode,
139 [
140 cols.a_pointer.into(),
141 cols.b_pointer.into(),
142 cols.c_pointer.into(),
143 native_as.clone(),
144 native_as.clone(),
145 ],
146 cols.from_state,
147 AB::F::from_canonical_usize(timestamp_delta),
148 (DEFAULT_PC_STEP, ctx.to_pc),
149 )
150 .eval(builder, ctx.instruction.is_valid);
151 }
152
153 fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
154 let cols: &NativeVectorizedAdapterCols<_, N> = local.borrow();
155 cols.from_state.pc
156 }
157}
158
159impl<F: PrimeField32, const N: usize> VmAdapterChip<F> for NativeVectorizedAdapterChip<F, N> {
160 type ReadRecord = NativeVectorizedReadRecord<N>;
161 type WriteRecord = NativeVectorizedWriteRecord<N>;
162 type Air = NativeVectorizedAdapterAir<N>;
163 type Interface = BasicAdapterInterface<F, MinimalInstruction<F>, 2, 1, N, N>;
164
165 fn preprocess(
166 &mut self,
167 memory: &mut MemoryController<F>,
168 instruction: &Instruction<F>,
169 ) -> Result<(
170 <Self::Interface as VmAdapterInterface<F>>::Reads,
171 Self::ReadRecord,
172 )> {
173 let Instruction { b, c, d, e, .. } = *instruction;
174
175 let y_val = memory.read::<N>(d, b);
176 let z_val = memory.read::<N>(e, c);
177
178 Ok((
179 [y_val.1, z_val.1],
180 Self::ReadRecord {
181 b: y_val.0,
182 c: z_val.0,
183 },
184 ))
185 }
186
187 fn postprocess(
188 &mut self,
189 memory: &mut MemoryController<F>,
190 instruction: &Instruction<F>,
191 from_state: ExecutionState<u32>,
192 output: AdapterRuntimeContext<F, Self::Interface>,
193 _read_record: &Self::ReadRecord,
194 ) -> Result<(ExecutionState<u32>, Self::WriteRecord)> {
195 let Instruction { a, d, .. } = *instruction;
196 let (a_val, _) = memory.write(d, a, output.writes[0]);
197
198 Ok((
199 ExecutionState {
200 pc: output.to_pc.unwrap_or(from_state.pc + DEFAULT_PC_STEP),
201 timestamp: memory.timestamp(),
202 },
203 Self::WriteRecord {
204 from_state,
205 a: a_val,
206 },
207 ))
208 }
209
210 fn generate_trace_row(
211 &self,
212 row_slice: &mut [F],
213 read_record: Self::ReadRecord,
214 write_record: Self::WriteRecord,
215 memory: &OfflineMemory<F>,
216 ) {
217 let aux_cols_factory = memory.aux_cols_factory();
218 let row_slice: &mut NativeVectorizedAdapterCols<_, N> = row_slice.borrow_mut();
219
220 let b_record = memory.record_by_id(read_record.b);
221 let c_record = memory.record_by_id(read_record.c);
222 let a_record = memory.record_by_id(write_record.a);
223 row_slice.from_state = write_record.from_state.map(F::from_canonical_u32);
224 row_slice.a_pointer = a_record.pointer;
225 row_slice.b_pointer = b_record.pointer;
226 row_slice.c_pointer = c_record.pointer;
227 aux_cols_factory.generate_read_aux(b_record, &mut row_slice.reads_aux[0]);
228 aux_cols_factory.generate_read_aux(c_record, &mut row_slice.reads_aux[1]);
229 aux_cols_factory.generate_write_aux(a_record, &mut row_slice.writes_aux[0]);
230 }
231
232 fn air(&self) -> &Self::Air {
233 &self.air
234 }
235}