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    #[cfg(not(feature = "tco"))]
67    fn pre_compute<Ctx: ExecutionCtxTrait>(
68        &self,
69        _pc: u32,
70        inst: &Instruction<F>,
71        data: &mut [u8],
72    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError> {
73        let data: &mut JalLuiPreCompute = data.borrow_mut();
74        let (is_jal, enabled) = self.pre_compute_impl(inst, data)?;
75        dispatch!(execute_e1_handler, is_jal, enabled)
76    }
77
78    #[cfg(feature = "tco")]
79    fn handler<Ctx>(
80        &self,
81        _pc: u32,
82        inst: &Instruction<F>,
83        data: &mut [u8],
84    ) -> Result<Handler<F, Ctx>, StaticProgramError>
85    where
86        Ctx: ExecutionCtxTrait,
87    {
88        let data: &mut JalLuiPreCompute = data.borrow_mut();
89        let (is_jal, enabled) = self.pre_compute_impl(inst, data)?;
90        dispatch!(execute_e1_handler, is_jal, enabled)
91    }
92}
93
94impl<F, A> MeteredExecutor<F> for Rv32JalLuiExecutor<A>
95where
96    F: PrimeField32,
97{
98    fn metered_pre_compute_size(&self) -> usize {
99        size_of::<E2PreCompute<JalLuiPreCompute>>()
100    }
101
102    #[cfg(not(feature = "tco"))]
103    fn metered_pre_compute<Ctx>(
104        &self,
105        chip_idx: usize,
106        _pc: u32,
107        inst: &Instruction<F>,
108        data: &mut [u8],
109    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
110    where
111        Ctx: MeteredExecutionCtxTrait,
112    {
113        let data: &mut E2PreCompute<JalLuiPreCompute> = data.borrow_mut();
114        data.chip_idx = chip_idx as u32;
115        let (is_jal, enabled) = self.pre_compute_impl(inst, &mut data.data)?;
116        dispatch!(execute_e2_handler, is_jal, enabled)
117    }
118
119    #[cfg(feature = "tco")]
120    fn metered_handler<Ctx>(
121        &self,
122        chip_idx: usize,
123        _pc: u32,
124        inst: &Instruction<F>,
125        data: &mut [u8],
126    ) -> Result<Handler<F, Ctx>, StaticProgramError>
127    where
128        Ctx: MeteredExecutionCtxTrait,
129    {
130        let data: &mut E2PreCompute<JalLuiPreCompute> = data.borrow_mut();
131        data.chip_idx = chip_idx as u32;
132        let (is_jal, enabled) = self.pre_compute_impl(inst, &mut data.data)?;
133        dispatch!(execute_e2_handler, is_jal, enabled)
134    }
135}
136
137#[inline(always)]
138unsafe fn execute_e12_impl<
139    F: PrimeField32,
140    CTX: ExecutionCtxTrait,
141    const IS_JAL: bool,
142    const ENABLED: bool,
143>(
144    pre_compute: &JalLuiPreCompute,
145    instret: &mut u64,
146    pc: &mut u32,
147    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
148) {
149    let JalLuiPreCompute { a, signed_imm } = *pre_compute;
150
151    let rd = if IS_JAL {
152        let rd_data = (*pc + DEFAULT_PC_STEP).to_le_bytes();
153        let next_pc = *pc as i32 + signed_imm;
154        debug_assert!(next_pc >= 0);
155        *pc = next_pc as u32;
156        rd_data
157    } else {
158        let imm = signed_imm as u32;
159        let rd = imm << 12;
160        *pc += DEFAULT_PC_STEP;
161        rd.to_le_bytes()
162    };
163
164    if ENABLED {
165        exec_state.vm_write(RV32_REGISTER_AS, a as u32, &rd);
166    }
167
168    *instret += 1;
169}
170
171#[create_handler]
172#[inline(always)]
173unsafe fn execute_e1_impl<
174    F: PrimeField32,
175    CTX: ExecutionCtxTrait,
176    const IS_JAL: bool,
177    const ENABLED: bool,
178>(
179    pre_compute: &[u8],
180    instret: &mut u64,
181    pc: &mut u32,
182    _instret_end: u64,
183    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
184) {
185    let pre_compute: &JalLuiPreCompute = pre_compute.borrow();
186    execute_e12_impl::<F, CTX, IS_JAL, ENABLED>(pre_compute, instret, pc, exec_state);
187}
188
189#[create_handler]
190#[inline(always)]
191unsafe fn execute_e2_impl<
192    F: PrimeField32,
193    CTX: MeteredExecutionCtxTrait,
194    const IS_JAL: bool,
195    const ENABLED: bool,
196>(
197    pre_compute: &[u8],
198    instret: &mut u64,
199    pc: &mut u32,
200    _arg: u64,
201    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
202) {
203    let pre_compute: &E2PreCompute<JalLuiPreCompute> = pre_compute.borrow();
204    exec_state
205        .ctx
206        .on_height_change(pre_compute.chip_idx as usize, 1);
207    execute_e12_impl::<F, CTX, IS_JAL, ENABLED>(&pre_compute.data, instret, pc, exec_state);
208}