openvm_rv32im_circuit/jalr/
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,
10    program::{DEFAULT_PC_STEP, PC_BITS},
11    riscv::RV32_REGISTER_AS,
12};
13use openvm_stark_backend::p3_field::PrimeField32;
14
15use super::core::Rv32JalrExecutor;
16
17#[derive(AlignedBytesBorrow, Clone)]
18#[repr(C)]
19struct JalrPreCompute {
20    imm_extended: u32,
21    a: u8,
22    b: u8,
23}
24
25impl<A> Rv32JalrExecutor<A> {
26    /// Return true if enabled.
27    fn pre_compute_impl<F: PrimeField32>(
28        &self,
29        pc: u32,
30        inst: &Instruction<F>,
31        data: &mut JalrPreCompute,
32    ) -> Result<bool, StaticProgramError> {
33        let imm_extended = inst.c.as_canonical_u32() + inst.g.as_canonical_u32() * 0xffff0000;
34        if inst.d.as_canonical_u32() != RV32_REGISTER_AS {
35            return Err(StaticProgramError::InvalidInstruction(pc));
36        }
37        *data = JalrPreCompute {
38            imm_extended,
39            a: inst.a.as_canonical_u32() as u8,
40            b: inst.b.as_canonical_u32() as u8,
41        };
42        let enabled = !inst.f.is_zero();
43        Ok(enabled)
44    }
45}
46
47macro_rules! dispatch {
48    ($execute_impl:ident, $enabled:ident) => {
49        if $enabled {
50            Ok($execute_impl::<_, _, true>)
51        } else {
52            Ok($execute_impl::<_, _, false>)
53        }
54    };
55}
56
57impl<F, A> Executor<F> for Rv32JalrExecutor<A>
58where
59    F: PrimeField32,
60{
61    #[inline(always)]
62    fn pre_compute_size(&self) -> usize {
63        size_of::<JalrPreCompute>()
64    }
65    #[inline(always)]
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 JalrPreCompute = data.borrow_mut();
73        let enabled = self.pre_compute_impl(pc, inst, data)?;
74        dispatch!(execute_e1_impl, 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 JalrPreCompute = data.borrow_mut();
88        let enabled = self.pre_compute_impl(pc, inst, data)?;
89        dispatch!(execute_e1_tco_handler, enabled)
90    }
91}
92
93impl<F, A> MeteredExecutor<F> for Rv32JalrExecutor<A>
94where
95    F: PrimeField32,
96{
97    fn metered_pre_compute_size(&self) -> usize {
98        size_of::<E2PreCompute<JalrPreCompute>>()
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<JalrPreCompute> = data.borrow_mut();
112        data.chip_idx = chip_idx as u32;
113        let enabled = self.pre_compute_impl(pc, inst, &mut data.data)?;
114        dispatch!(execute_e2_impl, 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<JalrPreCompute> = data.borrow_mut();
129        data.chip_idx = chip_idx as u32;
130        let enabled = self.pre_compute_impl(pc, inst, &mut data.data)?;
131        dispatch!(execute_e2_tco_handler, enabled)
132    }
133}
134
135#[inline(always)]
136unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const ENABLED: bool>(
137    pre_compute: &JalrPreCompute,
138    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
139) {
140    let rs1 = vm_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.b as u32);
141    let rs1 = u32::from_le_bytes(rs1);
142    let to_pc = rs1.wrapping_add(pre_compute.imm_extended);
143    let to_pc = to_pc - (to_pc & 1);
144    debug_assert!(to_pc < (1 << PC_BITS));
145    let rd = (vm_state.pc + DEFAULT_PC_STEP).to_le_bytes();
146
147    if ENABLED {
148        vm_state.vm_write(RV32_REGISTER_AS, pre_compute.a as u32, &rd);
149    }
150
151    vm_state.pc = to_pc;
152    vm_state.instret += 1;
153}
154
155#[create_tco_handler]
156unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const ENABLED: bool>(
157    pre_compute: &[u8],
158    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
159) {
160    let pre_compute: &JalrPreCompute = pre_compute.borrow();
161    execute_e12_impl::<F, CTX, ENABLED>(pre_compute, vm_state);
162}
163
164#[create_tco_handler]
165unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait, const ENABLED: bool>(
166    pre_compute: &[u8],
167    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
168) {
169    let pre_compute: &E2PreCompute<JalrPreCompute> = pre_compute.borrow();
170    vm_state
171        .ctx
172        .on_height_change(pre_compute.chip_idx as usize, 1);
173    execute_e12_impl::<F, CTX, ENABLED>(&pre_compute.data, vm_state);
174}
openvm_rv32im_circuit/jalr/execution.rs

openvm_rv32im_circuit/jalr/
execution.rs
