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}