openvm_native_circuit/adapters/
branch_native_adapter.rs

Help
1use std::{
2    borrow::{Borrow, BorrowMut},
3    marker::PhantomData,
4};
5
6use openvm_circuit::{
7    arch::{
8        AdapterAirContext, AdapterRuntimeContext, BasicAdapterInterface, ExecutionBridge,
9        ExecutionBus, ExecutionState, ImmInstruction, Result, VmAdapterAir, VmAdapterChip,
10        VmAdapterInterface,
11    },
12    system::{
13        memory::{
14            offline_checker::{MemoryBridge, MemoryReadOrImmediateAuxCols},
15            MemoryAddress, MemoryController, OfflineMemory,
16        },
17        native_adapter::NativeReadRecord,
18        program::ProgramBus,
19    },
20};
21use openvm_circuit_primitives_derive::AlignedBorrow;
22use openvm_instructions::{instruction::Instruction, program::DEFAULT_PC_STEP};
23use openvm_native_compiler::conversion::AS;
24use openvm_stark_backend::{
25    interaction::InteractionBuilder,
26    p3_air::BaseAir,
27    p3_field::{Field, FieldAlgebra, PrimeField32},
28};
29
30#[derive(Debug)]
31pub struct BranchNativeAdapterChip<F: Field> {
32    pub air: BranchNativeAdapterAir,
33    _marker: PhantomData<F>,
34}
35
36impl<F: PrimeField32> BranchNativeAdapterChip<F> {
37    pub fn new(
38        execution_bus: ExecutionBus,
39        program_bus: ProgramBus,
40        memory_bridge: MemoryBridge,
41    ) -> Self {
42        Self {
43            air: BranchNativeAdapterAir {
44                execution_bridge: ExecutionBridge::new(execution_bus, program_bus),
45                memory_bridge,
46            },
47            _marker: PhantomData,
48        }
49    }
50}
51
52#[repr(C)]
53#[derive(AlignedBorrow)]
54pub struct BranchNativeAdapterReadCols<T> {
55    pub address: MemoryAddress<T, T>,
56    pub read_aux: MemoryReadOrImmediateAuxCols<T>,
57}
58
59#[repr(C)]
60#[derive(AlignedBorrow)]
61pub struct BranchNativeAdapterCols<T> {
62    pub from_state: ExecutionState<T>,
63    pub reads_aux: [BranchNativeAdapterReadCols<T>; 2],
64}
65
66#[derive(Clone, Copy, Debug, derive_new::new)]
67pub struct BranchNativeAdapterAir {
68    pub(super) execution_bridge: ExecutionBridge,
69    pub(super) memory_bridge: MemoryBridge,
70}
71
72impl<F: Field> BaseAir<F> for BranchNativeAdapterAir {
73    fn width(&self) -> usize {
74        BranchNativeAdapterCols::<F>::width()
75    }
76}
77
78impl<AB: InteractionBuilder> VmAdapterAir<AB> for BranchNativeAdapterAir {
79    type Interface = BasicAdapterInterface<AB::Expr, ImmInstruction<AB::Expr>, 2, 0, 1, 1>;
80
81    fn eval(
82        &self,
83        builder: &mut AB,
84        local: &[AB::Var],
85        ctx: AdapterAirContext<AB::Expr, Self::Interface>,
86    ) {
87        let cols: &BranchNativeAdapterCols<_> = local.borrow();
88        let timestamp = cols.from_state.timestamp;
89        let mut timestamp_delta = 0usize;
90        let mut timestamp_pp = || {
91            timestamp_delta += 1;
92            timestamp + AB::F::from_canonical_usize(timestamp_delta - 1)
93        };
94
95        // check that d and e are in {0, 4}
96        let d = cols.reads_aux[0].address.address_space;
97        let e = cols.reads_aux[1].address.address_space;
98        builder.assert_eq(
99            d * (d - AB::F::from_canonical_u32(AS::Native as u32)),
100            AB::F::ZERO,
101        );
102        builder.assert_eq(
103            e * (e - AB::F::from_canonical_u32(AS::Native as u32)),
104            AB::F::ZERO,
105        );
106
107        self.memory_bridge
108            .read_or_immediate(
109                cols.reads_aux[0].address,
110                ctx.reads[0][0].clone(),
111                timestamp_pp(),
112                &cols.reads_aux[0].read_aux,
113            )
114            .eval(builder, ctx.instruction.is_valid.clone());
115
116        self.memory_bridge
117            .read_or_immediate(
118                cols.reads_aux[1].address,
119                ctx.reads[1][0].clone(),
120                timestamp_pp(),
121                &cols.reads_aux[1].read_aux,
122            )
123            .eval(builder, ctx.instruction.is_valid.clone());
124
125        self.execution_bridge
126            .execute_and_increment_or_set_pc(
127                ctx.instruction.opcode,
128                [
129                    cols.reads_aux[0].address.pointer.into(),
130                    cols.reads_aux[1].address.pointer.into(),
131                    ctx.instruction.immediate,
132                    cols.reads_aux[0].address.address_space.into(),
133                    cols.reads_aux[1].address.address_space.into(),
134                ],
135                cols.from_state,
136                AB::F::from_canonical_usize(timestamp_delta),
137                (DEFAULT_PC_STEP, ctx.to_pc),
138            )
139            .eval(builder, ctx.instruction.is_valid);
140    }
141
142    fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
143        let cols: &BranchNativeAdapterCols<_> = local.borrow();
144        cols.from_state.pc
145    }
146}
147
148impl<F: PrimeField32> VmAdapterChip<F> for BranchNativeAdapterChip<F> {
149    type ReadRecord = NativeReadRecord<F, 2>;
150    type WriteRecord = ExecutionState<u32>;
151    type Air = BranchNativeAdapterAir;
152    type Interface = BasicAdapterInterface<F, ImmInstruction<F>, 2, 0, 1, 1>;
153
154    fn preprocess(
155        &mut self,
156        memory: &mut MemoryController<F>,
157        instruction: &Instruction<F>,
158    ) -> Result<(
159        <Self::Interface as VmAdapterInterface<F>>::Reads,
160        Self::ReadRecord,
161    )> {
162        let Instruction { a, b, d, e, .. } = *instruction;
163
164        let reads = vec![memory.read::<1>(d, a), memory.read::<1>(e, b)];
165        let i_reads: [_; 2] = std::array::from_fn(|i| reads[i].1);
166
167        Ok((
168            i_reads,
169            Self::ReadRecord {
170                reads: reads.try_into().unwrap(),
171            },
172        ))
173    }
174
175    fn postprocess(
176        &mut self,
177        memory: &mut MemoryController<F>,
178        _instruction: &Instruction<F>,
179        from_state: ExecutionState<u32>,
180        output: AdapterRuntimeContext<F, Self::Interface>,
181        _read_record: &Self::ReadRecord,
182    ) -> Result<(ExecutionState<u32>, Self::WriteRecord)> {
183        Ok((
184            ExecutionState {
185                pc: output.to_pc.unwrap_or(from_state.pc + DEFAULT_PC_STEP),
186                timestamp: memory.timestamp(),
187            },
188            from_state,
189        ))
190    }
191
192    fn generate_trace_row(
193        &self,
194        row_slice: &mut [F],
195        read_record: Self::ReadRecord,
196        write_record: Self::WriteRecord,
197        memory: &OfflineMemory<F>,
198    ) {
199        let row_slice: &mut BranchNativeAdapterCols<_> = row_slice.borrow_mut();
200        let aux_cols_factory = memory.aux_cols_factory();
201
202        row_slice.from_state = write_record.map(F::from_canonical_u32);
203        for (i, x) in read_record.reads.iter().enumerate() {
204            let read = memory.record_by_id(x.0);
205
206            row_slice.reads_aux[i].address = MemoryAddress::new(read.address_space, read.pointer);
207            aux_cols_factory
208                .generate_read_or_immediate_aux(read, &mut row_slice.reads_aux[i].read_aux);
209        }
210    }
211
212    fn air(&self) -> &Self::Air {
213        &self.air
214    }
215}
openvm_native_circuit/adapters/branch_native_adapter.rs

openvm_native_circuit/adapters/
branch_native_adapter.rs
