openvm_rv32im_circuit/jal_lui/
execution.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    mem::size_of,
4};
5
6use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
7use openvm_circuit_primitives_derive::AlignedBytesBorrow;
8use openvm_instructions::{
9    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS, LocalOpcode,
10};
11use openvm_rv32im_transpiler::Rv32JalLuiOpcode::{self, JAL};
12use openvm_stark_backend::p3_field::PrimeField32;
13
14use super::core::{get_signed_imm, Rv32JalLuiExecutor};
15
16#[derive(AlignedBytesBorrow, Clone)]
17#[repr(C)]
18struct JalLuiPreCompute {
19    signed_imm: i32,
20    a: u8,
21}
22
23impl<A> Rv32JalLuiExecutor<A> {
24    /// Return (IS_JAL, ENABLED)
25    #[inline(always)]
26    fn pre_compute_impl<F: PrimeField32>(
27        &self,
28        inst: &Instruction<F>,
29        data: &mut JalLuiPreCompute,
30    ) -> Result<(bool, bool), StaticProgramError> {
31        let local_opcode = Rv32JalLuiOpcode::from_usize(
32            inst.opcode.local_opcode_idx(Rv32JalLuiOpcode::CLASS_OFFSET),
33        );
34        let is_jal = local_opcode == JAL;
35        let signed_imm = get_signed_imm(is_jal, inst.c);
36
37        *data = JalLuiPreCompute {
38            signed_imm,
39            a: inst.a.as_canonical_u32() as u8,
40        };
41        let enabled = !inst.f.is_zero();
42        Ok((is_jal, enabled))
43    }
44}
45
46macro_rules! dispatch {
47    ($execute_impl:ident, $is_jal:ident, $enabled:ident) => {
48        match ($is_jal, $enabled) {
49            (true, true) => Ok($execute_impl::<_, _, true, true>),
50            (true, false) => Ok($execute_impl::<_, _, true, false>),
51            (false, true) => Ok($execute_impl::<_, _, false, true>),
52            (false, false) => Ok($execute_impl::<_, _, false, false>),
53        }
54    };
55}
56
57impl<F, A> Executor<F> for Rv32JalLuiExecutor<A>
58where
59    F: PrimeField32,
60{
61    #[inline(always)]
62    fn pre_compute_size(&self) -> usize {
63        size_of::<JalLuiPreCompute>()
64    }
65
66    fn pre_compute<Ctx: ExecutionCtxTrait>(
67        &self,
68        _pc: u32,
69        inst: &Instruction<F>,
70        data: &mut [u8],
71    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError> {
72        let data: &mut JalLuiPreCompute = data.borrow_mut();
73        let (is_jal, enabled) = self.pre_compute_impl(inst, data)?;
74        dispatch!(execute_e1_impl, is_jal, enabled)
75    }
76
77    #[cfg(feature = "tco")]
78    fn handler<Ctx>(
79        &self,
80        _pc: u32,
81        inst: &Instruction<F>,
82        data: &mut [u8],
83    ) -> Result<Handler<F, Ctx>, StaticProgramError>
84    where
85        Ctx: ExecutionCtxTrait,
86    {
87        let data: &mut JalLuiPreCompute = data.borrow_mut();
88        let (is_jal, enabled) = self.pre_compute_impl(inst, data)?;
89        dispatch!(execute_e1_tco_handler, is_jal, enabled)
90    }
91}
92
93impl<F, A> MeteredExecutor<F> for Rv32JalLuiExecutor<A>
94where
95    F: PrimeField32,
96{
97    fn metered_pre_compute_size(&self) -> usize {
98        size_of::<E2PreCompute<JalLuiPreCompute>>()
99    }
100
101    fn metered_pre_compute<Ctx>(
102        &self,
103        chip_idx: usize,
104        _pc: u32,
105        inst: &Instruction<F>,
106        data: &mut [u8],
107    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
108    where
109        Ctx: MeteredExecutionCtxTrait,
110    {
111        let data: &mut E2PreCompute<JalLuiPreCompute> = data.borrow_mut();
112        data.chip_idx = chip_idx as u32;
113        let (is_jal, enabled) = self.pre_compute_impl(inst, &mut data.data)?;
114        dispatch!(execute_e2_impl, is_jal, enabled)
115    }
116
117    #[cfg(feature = "tco")]
118    fn metered_handler<Ctx>(
119        &self,
120        chip_idx: usize,
121        _pc: u32,
122        inst: &Instruction<F>,
123        data: &mut [u8],
124    ) -> Result<Handler<F, Ctx>, StaticProgramError>
125    where
126        Ctx: MeteredExecutionCtxTrait,
127    {
128        let data: &mut E2PreCompute<JalLuiPreCompute> = data.borrow_mut();
129        data.chip_idx = chip_idx as u32;
130        let (is_jal, enabled) = self.pre_compute_impl(inst, &mut data.data)?;
131        dispatch!(execute_e2_tco_handler, is_jal, enabled)
132    }
133}
134
135unsafe fn execute_e12_impl<
136    F: PrimeField32,
137    CTX: ExecutionCtxTrait,
138    const IS_JAL: bool,
139    const ENABLED: bool,
140>(
141    pre_compute: &JalLuiPreCompute,
142    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
143) {
144    let JalLuiPreCompute { a, signed_imm } = *pre_compute;
145
146    let rd = if IS_JAL {
147        let rd_data = (vm_state.pc + DEFAULT_PC_STEP).to_le_bytes();
148        let next_pc = vm_state.pc as i32 + signed_imm;
149        debug_assert!(next_pc >= 0);
150        vm_state.pc = next_pc as u32;
151        rd_data
152    } else {
153        let imm = signed_imm as u32;
154        let rd = imm << 12;
155        vm_state.pc += DEFAULT_PC_STEP;
156        rd.to_le_bytes()
157    };
158
159    if ENABLED {
160        vm_state.vm_write(RV32_REGISTER_AS, a as u32, &rd);
161    }
162
163    vm_state.instret += 1;
164}
165
166#[create_tco_handler]
167unsafe fn execute_e1_impl<
168    F: PrimeField32,
169    CTX: ExecutionCtxTrait,
170    const IS_JAL: bool,
171    const ENABLED: bool,
172>(
173    pre_compute: &[u8],
174    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
175) {
176    let pre_compute: &JalLuiPreCompute = pre_compute.borrow();
177    execute_e12_impl::<F, CTX, IS_JAL, ENABLED>(pre_compute, vm_state);
178}
179
180#[create_tco_handler]
181unsafe fn execute_e2_impl<
182    F: PrimeField32,
183    CTX: MeteredExecutionCtxTrait,
184    const IS_JAL: bool,
185    const ENABLED: bool,
186>(
187    pre_compute: &[u8],
188    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
189) {
190    let pre_compute: &E2PreCompute<JalLuiPreCompute> = pre_compute.borrow();
191    vm_state
192        .ctx
193        .on_height_change(pre_compute.chip_idx as usize, 1);
194    execute_e12_impl::<F, CTX, IS_JAL, ENABLED>(&pre_compute.data, vm_state);
195}