openvm_rv32im_circuit/adapters/
branch.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{
4    arch::{
5        get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller,
6        BasicAdapterInterface, ExecutionBridge, ExecutionState, ImmInstruction, VmAdapterAir,
7    },
8    system::memory::{
9        offline_checker::{MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord},
10        online::TracingMemory,
11        MemoryAddress, MemoryAuxColsFactory,
12    },
13};
14use openvm_circuit_primitives::{
15    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
16};
17use openvm_circuit_primitives_derive::AlignedBorrow;
18use openvm_instructions::{
19    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS,
20};
21use openvm_stark_backend::{
22    interaction::InteractionBuilder,
23    p3_air::BaseAir,
24    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
25};
26
27use super::RV32_REGISTER_NUM_LIMBS;
28use crate::adapters::tracing_read;
29
30#[repr(C)]
31#[derive(AlignedBorrow, StructReflection)]
32pub struct Rv32BranchAdapterCols<T> {
33    pub from_state: ExecutionState<T>,
34    pub rs1_ptr: T,
35    pub rs2_ptr: T,
36    pub reads_aux: [MemoryReadAuxCols<T>; 2],
37}
38
39#[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)]
40#[columns_via(Rv32BranchAdapterCols<u8>)]
41pub struct Rv32BranchAdapterAir {
42    pub(super) execution_bridge: ExecutionBridge,
43    pub(super) memory_bridge: MemoryBridge,
44}
45
46impl<F: Field> BaseAir<F> for Rv32BranchAdapterAir {
47    fn width(&self) -> usize {
48        Rv32BranchAdapterCols::<F>::width()
49    }
50}
51
52impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32BranchAdapterAir {
53    type Interface =
54        BasicAdapterInterface<AB::Expr, ImmInstruction<AB::Expr>, 2, 0, RV32_REGISTER_NUM_LIMBS, 0>;
55
56    fn eval(
57        &self,
58        builder: &mut AB,
59        local: &[AB::Var],
60        ctx: AdapterAirContext<AB::Expr, Self::Interface>,
61    ) {
62        let local: &Rv32BranchAdapterCols<_> = local.borrow();
63        let timestamp = local.from_state.timestamp;
64        let mut timestamp_delta: usize = 0;
65        let mut timestamp_pp = || {
66            timestamp_delta += 1;
67            timestamp + AB::F::from_usize(timestamp_delta - 1)
68        };
69
70        self.memory_bridge
71            .read(
72                MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), local.rs1_ptr),
73                ctx.reads[0].clone(),
74                timestamp_pp(),
75                &local.reads_aux[0],
76            )
77            .eval(builder, ctx.instruction.is_valid.clone());
78
79        self.memory_bridge
80            .read(
81                MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), local.rs2_ptr),
82                ctx.reads[1].clone(),
83                timestamp_pp(),
84                &local.reads_aux[1],
85            )
86            .eval(builder, ctx.instruction.is_valid.clone());
87
88        self.execution_bridge
89            .execute_and_increment_or_set_pc(
90                ctx.instruction.opcode,
91                [
92                    local.rs1_ptr.into(),
93                    local.rs2_ptr.into(),
94                    ctx.instruction.immediate,
95                    AB::Expr::from_u32(RV32_REGISTER_AS),
96                    AB::Expr::from_u32(RV32_REGISTER_AS),
97                ],
98                local.from_state,
99                AB::F::from_usize(timestamp_delta),
100                (DEFAULT_PC_STEP, ctx.to_pc),
101            )
102            .eval(builder, ctx.instruction.is_valid);
103    }
104
105    fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
106        let cols: &Rv32BranchAdapterCols<_> = local.borrow();
107        cols.from_state.pc
108    }
109}
110
111#[repr(C)]
112#[derive(AlignedBytesBorrow, Debug)]
113pub struct Rv32BranchAdapterRecord {
114    pub from_pc: u32,
115    pub from_timestamp: u32,
116    pub rs1_ptr: u32,
117    pub rs2_ptr: u32,
118    pub reads_aux: [MemoryReadAuxRecord; 2],
119}
120
121/// Reads instructions of the form OP a, b, c, d, e where if(\[a:4\]_d op \[b:4\]_e) pc += c.
122/// Operands d and e can only be 1.
123#[derive(Clone, Copy, derive_new::new)]
124pub struct Rv32BranchAdapterExecutor;
125
126#[derive(derive_new::new)]
127pub struct Rv32BranchAdapterFiller;
128
129impl<F> AdapterTraceExecutor<F> for Rv32BranchAdapterExecutor
130where
131    F: PrimeField32,
132{
133    const WIDTH: usize = size_of::<Rv32BranchAdapterCols<u8>>();
134    type ReadData = [[u8; RV32_REGISTER_NUM_LIMBS]; 2];
135    type WriteData = ();
136    type RecordMut<'a> = &'a mut Rv32BranchAdapterRecord;
137
138    #[inline(always)]
139    fn start(pc: u32, memory: &TracingMemory, record: &mut &mut Rv32BranchAdapterRecord) {
140        record.from_pc = pc;
141        record.from_timestamp = memory.timestamp;
142    }
143
144    #[inline(always)]
145    fn read(
146        &self,
147        memory: &mut TracingMemory,
148        instruction: &Instruction<F>,
149        record: &mut &mut Rv32BranchAdapterRecord,
150    ) -> Self::ReadData {
151        let &Instruction { a, b, d, e, .. } = instruction;
152
153        debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
154        debug_assert_eq!(e.as_canonical_u32(), RV32_REGISTER_AS);
155
156        record.rs1_ptr = a.as_canonical_u32();
157        let rs1 = tracing_read(
158            memory,
159            RV32_REGISTER_AS,
160            a.as_canonical_u32(),
161            &mut record.reads_aux[0].prev_timestamp,
162        );
163        record.rs2_ptr = b.as_canonical_u32();
164        let rs2 = tracing_read(
165            memory,
166            RV32_REGISTER_AS,
167            b.as_canonical_u32(),
168            &mut record.reads_aux[1].prev_timestamp,
169        );
170
171        [rs1, rs2]
172    }
173
174    #[inline(always)]
175    fn write(
176        &self,
177        _memory: &mut TracingMemory,
178        _instruction: &Instruction<F>,
179        _data: Self::WriteData,
180        _record: &mut Self::RecordMut<'_>,
181    ) {
182        // This function is intentionally left empty
183    }
184}
185
186impl<F: PrimeField32> AdapterTraceFiller<F> for Rv32BranchAdapterFiller {
187    const WIDTH: usize = size_of::<Rv32BranchAdapterCols<u8>>();
188
189    #[inline(always)]
190    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, mut adapter_row: &mut [F]) {
191        // SAFETY:
192        // - caller ensures `adapter_row` contains a valid record representation that was previously
193        //   written by the executor
194        // - get_record_from_slice correctly interprets the bytes as Rv32BranchAdapterRecord
195        let record: &Rv32BranchAdapterRecord =
196            unsafe { get_record_from_slice(&mut adapter_row, ()) };
197        let adapter_row: &mut Rv32BranchAdapterCols<F> = adapter_row.borrow_mut();
198
199        // We must assign in reverse
200        let timestamp = record.from_timestamp;
201
202        mem_helper.fill(
203            record.reads_aux[1].prev_timestamp,
204            timestamp + 1,
205            adapter_row.reads_aux[1].as_mut(),
206        );
207
208        mem_helper.fill(
209            record.reads_aux[0].prev_timestamp,
210            timestamp,
211            adapter_row.reads_aux[0].as_mut(),
212        );
213
214        adapter_row.from_state.pc = F::from_u32(record.from_pc);
215        adapter_row.from_state.timestamp = F::from_u32(record.from_timestamp);
216        adapter_row.rs1_ptr = F::from_u32(record.rs1_ptr);
217        adapter_row.rs2_ptr = F::from_u32(record.rs2_ptr);
218    }
219}
openvm_rv32im_circuit/adapters/branch.rs

openvm_rv32im_circuit/adapters/
branch.rs
