openvm_bigint_circuit/
branch_lt.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    mem::size_of,
4};
5
6use openvm_bigint_transpiler::Rv32BranchLessThan256Opcode;
7use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
8use openvm_circuit_primitives_derive::AlignedBytesBorrow;
9use openvm_instructions::{
10    instruction::Instruction,
11    program::DEFAULT_PC_STEP,
12    riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
13    LocalOpcode,
14};
15use openvm_rv32im_circuit::BranchLessThanExecutor;
16use openvm_rv32im_transpiler::BranchLessThanOpcode;
17use openvm_stark_backend::p3_field::PrimeField32;
18
19use crate::{
20    common::{i256_lt, read_int256, u256_lt},
21    BranchAdapterExecutor, Rv32BranchLessThan256Executor, INT256_NUM_LIMBS,
22};
23
24impl Rv32BranchLessThan256Executor {
25    pub fn new(adapter: BranchAdapterExecutor, offset: usize) -> Self {
26        Self(BranchLessThanExecutor::new(adapter, offset))
27    }
28}
29
30#[derive(AlignedBytesBorrow, Clone)]
31#[repr(C)]
32struct BranchLtPreCompute {
33    imm: isize,
34    a: u8,
35    b: u8,
36}
37
38macro_rules! dispatch {
39    ($execute_impl:ident, $local_opcode:ident) => {
40        Ok(match $local_opcode {
41            BranchLessThanOpcode::BLT => $execute_impl::<_, _, BltOp>,
42            BranchLessThanOpcode::BLTU => $execute_impl::<_, _, BltuOp>,
43            BranchLessThanOpcode::BGE => $execute_impl::<_, _, BgeOp>,
44            BranchLessThanOpcode::BGEU => $execute_impl::<_, _, BgeuOp>,
45        })
46    };
47}
48
49impl<F: PrimeField32> InterpreterExecutor<F> for Rv32BranchLessThan256Executor {
50    fn pre_compute_size(&self) -> usize {
51        size_of::<BranchLtPreCompute>()
52    }
53
54    #[cfg(not(feature = "tco"))]
55    fn pre_compute<Ctx>(
56        &self,
57        pc: u32,
58        inst: &Instruction<F>,
59        data: &mut [u8],
60    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
61    where
62        Ctx: ExecutionCtxTrait,
63    {
64        let data: &mut BranchLtPreCompute = data.borrow_mut();
65        let local_opcode = self.pre_compute_impl(pc, inst, data)?;
66        dispatch!(execute_e1_handler, local_opcode)
67    }
68
69    #[cfg(feature = "tco")]
70    fn handler<Ctx>(
71        &self,
72        pc: u32,
73        inst: &Instruction<F>,
74        data: &mut [u8],
75    ) -> Result<Handler<F, Ctx>, StaticProgramError>
76    where
77        Ctx: ExecutionCtxTrait,
78    {
79        let data: &mut BranchLtPreCompute = data.borrow_mut();
80        let local_opcode = self.pre_compute_impl(pc, inst, data)?;
81        dispatch!(execute_e1_handler, local_opcode)
82    }
83}
84
85#[cfg(feature = "aot")]
86impl<F: PrimeField32> AotExecutor<F> for Rv32BranchLessThan256Executor {}
87
88impl<F: PrimeField32> InterpreterMeteredExecutor<F> for Rv32BranchLessThan256Executor {
89    fn metered_pre_compute_size(&self) -> usize {
90        size_of::<E2PreCompute<BranchLtPreCompute>>()
91    }
92
93    #[cfg(not(feature = "tco"))]
94    fn metered_pre_compute<Ctx>(
95        &self,
96        chip_idx: usize,
97        pc: u32,
98        inst: &Instruction<F>,
99        data: &mut [u8],
100    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
101    where
102        Ctx: MeteredExecutionCtxTrait,
103    {
104        let data: &mut E2PreCompute<BranchLtPreCompute> = data.borrow_mut();
105        data.chip_idx = chip_idx as u32;
106        let local_opcode = self.pre_compute_impl(pc, inst, &mut data.data)?;
107        dispatch!(execute_e2_handler, local_opcode)
108    }
109
110    #[cfg(feature = "tco")]
111    fn metered_handler<Ctx>(
112        &self,
113        chip_idx: usize,
114        pc: u32,
115        inst: &Instruction<F>,
116        data: &mut [u8],
117    ) -> Result<Handler<F, Ctx>, StaticProgramError>
118    where
119        Ctx: MeteredExecutionCtxTrait,
120    {
121        let data: &mut E2PreCompute<BranchLtPreCompute> = data.borrow_mut();
122        data.chip_idx = chip_idx as u32;
123        let local_opcode = self.pre_compute_impl(pc, inst, &mut data.data)?;
124        dispatch!(execute_e2_handler, local_opcode)
125    }
126}
127
128#[cfg(feature = "aot")]
129impl<F: PrimeField32> AotMeteredExecutor<F> for Rv32BranchLessThan256Executor {}
130
131#[inline(always)]
132unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, OP: BranchLessThanOp>(
133    pre_compute: &BranchLtPreCompute,
134    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
135) {
136    let mut pc = exec_state.pc();
137    let rs1_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.a as u32);
138    let rs2_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.b as u32);
139    let rs1 = read_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rs1_ptr));
140    let rs2 = read_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rs2_ptr));
141    let cmp_result = OP::compute(rs1, rs2);
142    if cmp_result {
143        pc = (pc as isize + pre_compute.imm) as u32;
144    } else {
145        pc = pc.wrapping_add(DEFAULT_PC_STEP);
146    }
147    exec_state.set_pc(pc);
148}
149
150#[create_handler]
151#[inline(always)]
152unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait, OP: BranchLessThanOp>(
153    pre_compute: *const u8,
154    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
155) {
156    let pre_compute: &BranchLtPreCompute =
157        std::slice::from_raw_parts(pre_compute, size_of::<BranchLtPreCompute>()).borrow();
158    execute_e12_impl::<F, CTX, OP>(pre_compute, exec_state);
159}
160
161#[create_handler]
162#[inline(always)]
163unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait, OP: BranchLessThanOp>(
164    pre_compute: *const u8,
165    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
166) {
167    let pre_compute: &E2PreCompute<BranchLtPreCompute> =
168        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<BranchLtPreCompute>>())
169            .borrow();
170    exec_state
171        .ctx
172        .on_height_change(pre_compute.chip_idx as usize, 1);
173    execute_e12_impl::<F, CTX, OP>(&pre_compute.data, exec_state);
174}
175
176impl Rv32BranchLessThan256Executor {
177    fn pre_compute_impl<F: PrimeField32>(
178        &self,
179        pc: u32,
180        inst: &Instruction<F>,
181        data: &mut BranchLtPreCompute,
182    ) -> Result<BranchLessThanOpcode, StaticProgramError> {
183        let Instruction {
184            opcode,
185            a,
186            b,
187            c,
188            d,
189            e,
190            ..
191        } = inst;
192        let c = c.as_canonical_u32();
193        let imm = if F::ORDER_U32 - c < c {
194            -((F::ORDER_U32 - c) as isize)
195        } else {
196            c as isize
197        };
198        let e_u32 = e.as_canonical_u32();
199        if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 != RV32_MEMORY_AS {
200            return Err(StaticProgramError::InvalidInstruction(pc));
201        }
202        *data = BranchLtPreCompute {
203            imm,
204            a: a.as_canonical_u32() as u8,
205            b: b.as_canonical_u32() as u8,
206        };
207        let local_opcode = BranchLessThanOpcode::from_usize(
208            opcode.local_opcode_idx(Rv32BranchLessThan256Opcode::CLASS_OFFSET),
209        );
210        Ok(local_opcode)
211    }
212}
213
214trait BranchLessThanOp {
215    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> bool;
216}
217struct BltOp;
218struct BltuOp;
219struct BgeOp;
220struct BgeuOp;
221
222impl BranchLessThanOp for BltOp {
223    #[inline(always)]
224    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> bool {
225        i256_lt(rs1, rs2)
226    }
227}
228impl BranchLessThanOp for BltuOp {
229    #[inline(always)]
230    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> bool {
231        u256_lt(rs1, rs2)
232    }
233}
234impl BranchLessThanOp for BgeOp {
235    #[inline(always)]
236    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> bool {
237        !i256_lt(rs1, rs2)
238    }
239}
240impl BranchLessThanOp for BgeuOp {
241    #[inline(always)]
242    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> bool {
243        !u256_lt(rs1, rs2)
244    }
245}