openvm_rv32im_circuit/adapters/
jalr.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{
4    arch::{
5        get_record_from_slice, AdapterAirContext, AdapterTraceExecutor, AdapterTraceFiller,
6        BasicAdapterInterface, ExecutionBridge, ExecutionState, SignedImmInstruction, VmAdapterAir,
7    },
8    system::memory::{
9        offline_checker::{
10            MemoryBridge, MemoryReadAuxCols, MemoryReadAuxRecord, MemoryWriteAuxCols,
11            MemoryWriteBytesAuxRecord,
12        },
13        online::TracingMemory,
14        MemoryAddress, MemoryAuxColsFactory,
15    },
16};
17use openvm_circuit_primitives::{
18    utils::not, AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
19};
20use openvm_circuit_primitives_derive::AlignedBorrow;
21use openvm_instructions::{
22    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS,
23};
24use openvm_stark_backend::{
25    interaction::InteractionBuilder,
26    p3_air::{AirBuilder, BaseAir},
27    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
28};
29
30use super::RV32_REGISTER_NUM_LIMBS;
31use crate::adapters::{tracing_read, tracing_write};
32
33#[repr(C)]
34#[derive(Debug, Clone, AlignedBorrow, StructReflection)]
35pub struct Rv32JalrAdapterCols<T> {
36    pub from_state: ExecutionState<T>,
37    pub rs1_ptr: T,
38    pub rs1_aux_cols: MemoryReadAuxCols<T>,
39    pub rd_ptr: T,
40    pub rd_aux_cols: MemoryWriteAuxCols<T, RV32_REGISTER_NUM_LIMBS>,
41    /// Only writes if `needs_write`.
42    /// Sets `needs_write` to 0 iff `rd == x0`
43    pub needs_write: T,
44}
45
46#[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)]
47#[columns_via(Rv32JalrAdapterCols<u8>)]
48pub struct Rv32JalrAdapterAir {
49    pub(super) memory_bridge: MemoryBridge,
50    pub(super) execution_bridge: ExecutionBridge,
51}
52
53impl<F: Field> BaseAir<F> for Rv32JalrAdapterAir {
54    fn width(&self) -> usize {
55        Rv32JalrAdapterCols::<F>::width()
56    }
57}
58
59impl<AB: InteractionBuilder> VmAdapterAir<AB> for Rv32JalrAdapterAir {
60    type Interface = BasicAdapterInterface<
61        AB::Expr,
62        SignedImmInstruction<AB::Expr>,
63        1,
64        1,
65        RV32_REGISTER_NUM_LIMBS,
66        RV32_REGISTER_NUM_LIMBS,
67    >;
68
69    fn eval(
70        &self,
71        builder: &mut AB,
72        local: &[AB::Var],
73        ctx: AdapterAirContext<AB::Expr, Self::Interface>,
74    ) {
75        let local_cols: &Rv32JalrAdapterCols<AB::Var> = local.borrow();
76
77        let timestamp: AB::Var = local_cols.from_state.timestamp;
78        let mut timestamp_delta: usize = 0;
79        let mut timestamp_pp = || {
80            timestamp_delta += 1;
81            timestamp + AB::Expr::from_usize(timestamp_delta - 1)
82        };
83
84        let write_count = local_cols.needs_write;
85
86        builder.assert_bool(write_count);
87        builder
88            .when::<AB::Expr>(not(ctx.instruction.is_valid.clone()))
89            .assert_zero(write_count);
90
91        self.memory_bridge
92            .read(
93                MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), local_cols.rs1_ptr),
94                ctx.reads[0].clone(),
95                timestamp_pp(),
96                &local_cols.rs1_aux_cols,
97            )
98            .eval(builder, ctx.instruction.is_valid.clone());
99
100        self.memory_bridge
101            .write(
102                MemoryAddress::new(AB::F::from_u32(RV32_REGISTER_AS), local_cols.rd_ptr),
103                ctx.writes[0].clone(),
104                timestamp_pp(),
105                &local_cols.rd_aux_cols,
106            )
107            .eval(builder, write_count);
108
109        let to_pc = ctx
110            .to_pc
111            .unwrap_or(local_cols.from_state.pc + AB::F::from_u32(DEFAULT_PC_STEP));
112
113        // regardless of `needs_write`, must always execute instruction when `is_valid`.
114        self.execution_bridge
115            .execute(
116                ctx.instruction.opcode,
117                [
118                    local_cols.rd_ptr.into(),
119                    local_cols.rs1_ptr.into(),
120                    ctx.instruction.immediate,
121                    AB::Expr::from_u32(RV32_REGISTER_AS),
122                    AB::Expr::ZERO,
123                    write_count.into(),
124                    ctx.instruction.imm_sign,
125                ],
126                local_cols.from_state,
127                ExecutionState {
128                    pc: to_pc,
129                    timestamp: timestamp + AB::F::from_usize(timestamp_delta),
130                },
131            )
132            .eval(builder, ctx.instruction.is_valid);
133    }
134
135    fn get_from_pc(&self, local: &[AB::Var]) -> AB::Var {
136        let cols: &Rv32JalrAdapterCols<_> = local.borrow();
137        cols.from_state.pc
138    }
139}
140
141#[repr(C)]
142#[derive(AlignedBytesBorrow, Debug)]
143pub struct Rv32JalrAdapterRecord {
144    pub from_pc: u32,
145    pub from_timestamp: u32,
146
147    pub rs1_ptr: u32,
148    // Will use u32::MAX to indicate no write
149    pub rd_ptr: u32,
150
151    pub reads_aux: MemoryReadAuxRecord,
152    pub writes_aux: MemoryWriteBytesAuxRecord<RV32_REGISTER_NUM_LIMBS>,
153}
154
155// This adapter reads from [b:4]_d (rs1) and writes to [a:4]_d (rd)
156#[derive(Clone, Copy, derive_new::new)]
157pub struct Rv32JalrAdapterExecutor;
158
159#[derive(Clone, Copy, derive_new::new)]
160pub struct Rv32JalrAdapterFiller;
161
162impl<F> AdapterTraceExecutor<F> for Rv32JalrAdapterExecutor
163where
164    F: PrimeField32,
165{
166    const WIDTH: usize = size_of::<Rv32JalrAdapterCols<u8>>();
167    type ReadData = [u8; RV32_REGISTER_NUM_LIMBS];
168    type WriteData = [u8; RV32_REGISTER_NUM_LIMBS];
169    type RecordMut<'a> = &'a mut Rv32JalrAdapterRecord;
170
171    #[inline(always)]
172    fn start(pc: u32, memory: &TracingMemory, record: &mut Self::RecordMut<'_>) {
173        record.from_pc = pc;
174        record.from_timestamp = memory.timestamp;
175    }
176
177    #[inline(always)]
178    fn read(
179        &self,
180        memory: &mut TracingMemory,
181        instruction: &Instruction<F>,
182        record: &mut Self::RecordMut<'_>,
183    ) -> Self::ReadData {
184        let &Instruction { b, d, .. } = instruction;
185
186        debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
187
188        record.rs1_ptr = b.as_canonical_u32();
189        tracing_read(
190            memory,
191            RV32_REGISTER_AS,
192            b.as_canonical_u32(),
193            &mut record.reads_aux.prev_timestamp,
194        )
195    }
196
197    #[inline(always)]
198    fn write(
199        &self,
200        memory: &mut TracingMemory,
201        instruction: &Instruction<F>,
202        data: Self::WriteData,
203        record: &mut Self::RecordMut<'_>,
204    ) {
205        let &Instruction {
206            a, d, f: enabled, ..
207        } = instruction;
208
209        debug_assert_eq!(d.as_canonical_u32(), RV32_REGISTER_AS);
210
211        if enabled.is_one() {
212            record.rd_ptr = a.as_canonical_u32();
213
214            tracing_write(
215                memory,
216                RV32_REGISTER_AS,
217                a.as_canonical_u32(),
218                data,
219                &mut record.writes_aux.prev_timestamp,
220                &mut record.writes_aux.prev_data,
221            );
222        } else {
223            record.rd_ptr = u32::MAX;
224            memory.increment_timestamp();
225        }
226    }
227}
228
229impl<F: PrimeField32> AdapterTraceFiller<F> for Rv32JalrAdapterFiller {
230    const WIDTH: usize = size_of::<Rv32JalrAdapterCols<u8>>();
231
232    #[inline(always)]
233    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, mut adapter_row: &mut [F]) {
234        // SAFETY:
235        // - caller ensures `adapter_row` contains a valid record representation that was previously
236        //   written by the executor
237        // - get_record_from_slice correctly interprets the bytes as Rv32JalrAdapterRecord
238        let record: &Rv32JalrAdapterRecord = unsafe { get_record_from_slice(&mut adapter_row, ()) };
239        let adapter_row: &mut Rv32JalrAdapterCols<F> = adapter_row.borrow_mut();
240
241        // We must assign in reverse
242        adapter_row.needs_write = F::from_bool(record.rd_ptr != u32::MAX);
243
244        if record.rd_ptr != u32::MAX {
245            adapter_row
246                .rd_aux_cols
247                .set_prev_data(record.writes_aux.prev_data.map(F::from_u8));
248            mem_helper.fill(
249                record.writes_aux.prev_timestamp,
250                record.from_timestamp + 1,
251                adapter_row.rd_aux_cols.as_mut(),
252            );
253            adapter_row.rd_ptr = F::from_u32(record.rd_ptr);
254        } else {
255            adapter_row.rd_ptr = F::ZERO;
256        }
257
258        mem_helper.fill(
259            record.reads_aux.prev_timestamp,
260            record.from_timestamp,
261            adapter_row.rs1_aux_cols.as_mut(),
262        );
263        adapter_row.rs1_ptr = F::from_u32(record.rs1_ptr);
264        adapter_row.from_state.timestamp = F::from_u32(record.from_timestamp);
265        adapter_row.from_state.pc = F::from_u32(record.from_pc);
266    }
267}
openvm_rv32im_circuit/adapters/jalr.rs

openvm_rv32im_circuit/adapters/
jalr.rs
