openvm_rv32im_circuit/jal_lui/
core.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{
4    arch::*,
5    system::memory::{online::TracingMemory, MemoryAuxColsFactory},
6};
7use openvm_circuit_primitives::{
8    bitwise_op_lookup::{BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip},
9    AlignedBytesBorrow, ColumnsAir, StructReflection, StructReflectionHelper,
10};
11use openvm_circuit_primitives_derive::AlignedBorrow;
12use openvm_instructions::{
13    instruction::Instruction,
14    program::{DEFAULT_PC_STEP, PC_BITS},
15    LocalOpcode,
16};
17use openvm_rv32im_transpiler::Rv32JalLuiOpcode::{self, *};
18use openvm_stark_backend::{
19    interaction::InteractionBuilder,
20    p3_air::{AirBuilder, BaseAir},
21    p3_field::{Field, PrimeCharacteristicRing, PrimeField32},
22    BaseAirWithPublicValues,
23};
24
25use crate::adapters::{
26    Rv32CondRdWriteAdapterExecutor, Rv32CondRdWriteAdapterFiller, RV32_CELL_BITS,
27    RV32_REGISTER_NUM_LIMBS, RV_J_TYPE_IMM_BITS,
28};
29
30pub(super) const ADDITIONAL_BITS: u32 = 0b11000000;
31
32#[repr(C)]
33#[derive(Debug, Clone, AlignedBorrow, StructReflection)]
34pub struct Rv32JalLuiCoreCols<T> {
35    pub imm: T,
36    pub rd_data: [T; RV32_REGISTER_NUM_LIMBS],
37    pub is_jal: T,
38    pub is_lui: T,
39}
40
41#[derive(Debug, Clone, Copy, derive_new::new, ColumnsAir)]
42#[columns_via(Rv32JalLuiCoreCols<u8>)]
43pub struct Rv32JalLuiCoreAir {
44    pub bus: BitwiseOperationLookupBus,
45}
46
47impl<F: Field> BaseAir<F> for Rv32JalLuiCoreAir {
48    fn width(&self) -> usize {
49        Rv32JalLuiCoreCols::<F>::width()
50    }
51}
52
53impl<F: Field> BaseAirWithPublicValues<F> for Rv32JalLuiCoreAir {}
54
55impl<AB, I> VmCoreAir<AB, I> for Rv32JalLuiCoreAir
56where
57    AB: InteractionBuilder,
58    I: VmAdapterInterface<AB::Expr>,
59    I::Reads: From<[[AB::Expr; 0]; 0]>,
60    I::Writes: From<[[AB::Expr; RV32_REGISTER_NUM_LIMBS]; 1]>,
61    I::ProcessedInstruction: From<ImmInstruction<AB::Expr>>,
62{
63    fn eval(
64        &self,
65        builder: &mut AB,
66        local_core: &[AB::Var],
67        from_pc: AB::Var,
68    ) -> AdapterAirContext<AB::Expr, I> {
69        let cols: &Rv32JalLuiCoreCols<AB::Var> = (*local_core).borrow();
70        let Rv32JalLuiCoreCols::<AB::Var> {
71            imm,
72            rd_data: rd,
73            is_jal,
74            is_lui,
75        } = *cols;
76
77        builder.assert_bool(is_lui);
78        builder.assert_bool(is_jal);
79        let is_valid = is_lui + is_jal;
80        builder.assert_bool(is_valid.clone());
81        builder.when(is_lui).assert_zero(rd[0]);
82
83        for i in 0..RV32_REGISTER_NUM_LIMBS / 2 {
84            self.bus
85                .send_range(rd[i * 2], rd[i * 2 + 1])
86                .eval(builder, is_valid.clone());
87        }
88
89        // In case of JAL constrain that last limb has at most [last_limb_bits] bits
90
91        let last_limb_bits = PC_BITS - RV32_CELL_BITS * (RV32_REGISTER_NUM_LIMBS - 1);
92        let additional_bits = (last_limb_bits..RV32_CELL_BITS).fold(0, |acc, x| acc + (1 << x));
93        let additional_bits = AB::F::from_u32(additional_bits);
94        self.bus
95            .send_xor(rd[3], additional_bits, rd[3] + additional_bits)
96            .eval(builder, is_jal);
97
98        let intermed_val = rd
99            .iter()
100            .skip(1)
101            .enumerate()
102            .fold(AB::Expr::ZERO, |acc, (i, &val)| {
103                acc + val * AB::Expr::from_u32(1 << (i * RV32_CELL_BITS))
104            });
105
106        // Constrain that imm * 2^4 is the correct composition of intermed_val in case of LUI
107        builder.when(is_lui).assert_eq(
108            intermed_val.clone(),
109            imm * AB::F::from_u32(1 << (12 - RV32_CELL_BITS)),
110        );
111
112        let intermed_val = rd[0] + intermed_val * AB::Expr::from_u32(1 << RV32_CELL_BITS);
113        // Constrain that from_pc + DEFAULT_PC_STEP is the correct composition of intermed_val in
114        // case of JAL
115        builder
116            .when(is_jal)
117            .assert_eq(intermed_val, from_pc + AB::F::from_u32(DEFAULT_PC_STEP));
118
119        let to_pc = from_pc + is_lui * AB::F::from_u32(DEFAULT_PC_STEP) + is_jal * imm;
120
121        let expected_opcode = VmCoreAir::<AB, I>::expr_to_global_expr(
122            self,
123            is_lui * AB::F::from_u32(LUI as u32) + is_jal * AB::F::from_u32(JAL as u32),
124        );
125
126        AdapterAirContext {
127            to_pc: Some(to_pc),
128            reads: [].into(),
129            writes: [rd.map(|x| x.into())].into(),
130            instruction: ImmInstruction {
131                is_valid,
132                opcode: expected_opcode,
133                immediate: imm.into(),
134            }
135            .into(),
136        }
137    }
138
139    fn start_offset(&self) -> usize {
140        Rv32JalLuiOpcode::CLASS_OFFSET
141    }
142}
143
144#[repr(C)]
145#[derive(AlignedBytesBorrow, Debug)]
146pub struct Rv32JalLuiCoreRecord {
147    pub imm: u32,
148    pub rd_data: [u8; RV32_REGISTER_NUM_LIMBS],
149    pub is_jal: bool,
150}
151
152#[derive(Clone, Copy, derive_new::new)]
153pub struct Rv32JalLuiExecutor<A = Rv32CondRdWriteAdapterExecutor> {
154    pub adapter: A,
155}
156
157#[derive(Clone, derive_new::new)]
158pub struct Rv32JalLuiFiller<A = Rv32CondRdWriteAdapterFiller> {
159    adapter: A,
160    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<RV32_CELL_BITS>,
161}
162
163impl<F, A, RA> PreflightExecutor<F, RA> for Rv32JalLuiExecutor<A>
164where
165    F: PrimeField32,
166    A: 'static
167        + for<'a> AdapterTraceExecutor<F, ReadData = (), WriteData = [u8; RV32_REGISTER_NUM_LIMBS]>,
168    for<'buf> RA: RecordArena<
169        'buf,
170        EmptyAdapterCoreLayout<F, A>,
171        (A::RecordMut<'buf>, &'buf mut Rv32JalLuiCoreRecord),
172    >,
173{
174    fn get_opcode_name(&self, opcode: usize) -> String {
175        format!(
176            "{:?}",
177            Rv32JalLuiOpcode::from_usize(opcode - Rv32JalLuiOpcode::CLASS_OFFSET)
178        )
179    }
180
181    fn execute(
182        &self,
183        state: VmStateMut<F, TracingMemory, RA>,
184        instruction: &Instruction<F>,
185    ) -> Result<(), ExecutionError> {
186        let &Instruction { opcode, c: imm, .. } = instruction;
187
188        let (mut adapter_record, core_record) = state.ctx.alloc(EmptyAdapterCoreLayout::new());
189
190        A::start(*state.pc, state.memory, &mut adapter_record);
191
192        let is_jal = opcode.local_opcode_idx(Rv32JalLuiOpcode::CLASS_OFFSET) == JAL as usize;
193        let signed_imm = get_signed_imm(is_jal, imm);
194
195        let (to_pc, rd_data) = run_jal_lui(is_jal, *state.pc, signed_imm);
196
197        core_record.imm = imm.as_canonical_u32();
198        core_record.rd_data = rd_data;
199        core_record.is_jal = is_jal;
200
201        self.adapter
202            .write(state.memory, instruction, rd_data, &mut adapter_record);
203
204        *state.pc = to_pc;
205
206        Ok(())
207    }
208}
209
210impl<F, A> TraceFiller<F> for Rv32JalLuiFiller<A>
211where
212    F: PrimeField32,
213    A: 'static + AdapterTraceFiller<F>,
214{
215    fn fill_trace_row(&self, mem_helper: &MemoryAuxColsFactory<F>, row_slice: &mut [F]) {
216        // SAFETY: row_slice is guaranteed by the caller to have at least A::WIDTH +
217        // Rv32JalLuiCoreCols::width() elements
218        let (adapter_row, mut core_row) = unsafe { row_slice.split_at_mut_unchecked(A::WIDTH) };
219        self.adapter.fill_trace_row(mem_helper, adapter_row);
220        // SAFETY: core_row contains a valid Rv32JalLuiCoreRecord written by the executor
221        // during trace generation
222        let record: &Rv32JalLuiCoreRecord = unsafe { get_record_from_slice(&mut core_row, ()) };
223        let core_row: &mut Rv32JalLuiCoreCols<F> = core_row.borrow_mut();
224
225        for pair in record.rd_data.chunks_exact(2) {
226            self.bitwise_lookup_chip
227                .request_range(pair[0] as u32, pair[1] as u32);
228        }
229        if record.is_jal {
230            self.bitwise_lookup_chip
231                .request_xor(record.rd_data[3] as u32, ADDITIONAL_BITS);
232        }
233
234        // Writing in reverse order
235        core_row.is_lui = F::from_bool(!record.is_jal);
236        core_row.is_jal = F::from_bool(record.is_jal);
237        core_row.rd_data = record.rd_data.map(F::from_u8);
238        core_row.imm = F::from_u32(record.imm);
239    }
240}
241
242// returns the canonical signed representation of the immediate
243// `imm` can be "negative" as a field element
244pub(super) fn get_signed_imm<F: PrimeField32>(is_jal: bool, imm: F) -> i32 {
245    let imm_f = imm.as_canonical_u32();
246    if is_jal {
247        if imm_f < (1 << (RV_J_TYPE_IMM_BITS - 1)) {
248            imm_f as i32
249        } else {
250            let neg_imm_f = F::ORDER_U32 - imm_f;
251            debug_assert!(neg_imm_f < (1 << (RV_J_TYPE_IMM_BITS - 1)));
252            -(neg_imm_f as i32)
253        }
254    } else {
255        imm_f as i32
256    }
257}
258
259// returns (to_pc, rd_data)
260#[inline(always)]
261pub(super) fn run_jal_lui(is_jal: bool, pc: u32, imm: i32) -> (u32, [u8; RV32_REGISTER_NUM_LIMBS]) {
262    if is_jal {
263        let rd_data = (pc + DEFAULT_PC_STEP).to_le_bytes();
264        let next_pc = pc as i32 + imm;
265        debug_assert!(next_pc >= 0);
266        (next_pc as u32, rd_data)
267    } else {
268        let imm = imm as u32;
269        let rd = imm << 12;
270        (pc + DEFAULT_PC_STEP, rd.to_le_bytes())
271    }
272}
openvm_rv32im_circuit/jal_lui/core.rs

openvm_rv32im_circuit/jal_lui/
core.rs
