openvm_bigint_circuit/
shift.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    mem::size_of,
4};
5
6use openvm_bigint_transpiler::Rv32Shift256Opcode;
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::ShiftExecutor;
16use openvm_rv32im_transpiler::ShiftOpcode;
17use openvm_stark_backend::p3_field::PrimeField32;
18
19use crate::{
20    common::{bytes_to_u64_array, read_int256, u64_array_to_bytes, write_int256},
21    AluAdapterExecutor, Rv32Shift256Executor, INT256_NUM_LIMBS,
22};
23
24impl Rv32Shift256Executor {
25    pub fn new(adapter: AluAdapterExecutor, offset: usize) -> Self {
26        Self(ShiftExecutor::new(adapter, offset))
27    }
28}
29
30#[derive(AlignedBytesBorrow, Clone)]
31#[repr(C)]
32struct ShiftPreCompute {
33    a: u8,
34    b: u8,
35    c: u8,
36}
37
38macro_rules! dispatch {
39    ($execute_impl:ident, $local_opcode:ident) => {
40        Ok(match $local_opcode {
41            ShiftOpcode::SLL => $execute_impl::<_, _, SllOp>,
42            ShiftOpcode::SRA => $execute_impl::<_, _, SraOp>,
43            ShiftOpcode::SRL => $execute_impl::<_, _, SrlOp>,
44        })
45    };
46}
47
48impl<F: PrimeField32> InterpreterExecutor<F> for Rv32Shift256Executor {
49    fn pre_compute_size(&self) -> usize {
50        size_of::<ShiftPreCompute>()
51    }
52
53    #[cfg(not(feature = "tco"))]
54    fn pre_compute<Ctx>(
55        &self,
56        pc: u32,
57        inst: &Instruction<F>,
58        data: &mut [u8],
59    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
60    where
61        Ctx: ExecutionCtxTrait,
62    {
63        let data: &mut ShiftPreCompute = data.borrow_mut();
64        let local_opcode = self.pre_compute_impl(pc, inst, data)?;
65        dispatch!(execute_e1_handler, local_opcode)
66    }
67
68    #[cfg(feature = "tco")]
69    fn handler<Ctx>(
70        &self,
71        pc: u32,
72        inst: &Instruction<F>,
73        data: &mut [u8],
74    ) -> Result<Handler<F, Ctx>, StaticProgramError>
75    where
76        Ctx: ExecutionCtxTrait,
77    {
78        let data: &mut ShiftPreCompute = data.borrow_mut();
79        let local_opcode = self.pre_compute_impl(pc, inst, data)?;
80        dispatch!(execute_e1_handler, local_opcode)
81    }
82}
83
84#[cfg(feature = "aot")]
85impl<F: PrimeField32> AotExecutor<F> for Rv32Shift256Executor {}
86
87impl<F: PrimeField32> InterpreterMeteredExecutor<F> for Rv32Shift256Executor {
88    fn metered_pre_compute_size(&self) -> usize {
89        size_of::<E2PreCompute<ShiftPreCompute>>()
90    }
91
92    #[cfg(not(feature = "tco"))]
93    fn metered_pre_compute<Ctx>(
94        &self,
95        chip_idx: usize,
96        pc: u32,
97        inst: &Instruction<F>,
98        data: &mut [u8],
99    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
100    where
101        Ctx: MeteredExecutionCtxTrait,
102    {
103        let data: &mut E2PreCompute<ShiftPreCompute> = data.borrow_mut();
104        data.chip_idx = chip_idx as u32;
105        let local_opcode = self.pre_compute_impl(pc, inst, &mut data.data)?;
106        dispatch!(execute_e2_handler, local_opcode)
107    }
108
109    #[cfg(feature = "tco")]
110    fn metered_handler<Ctx>(
111        &self,
112        chip_idx: usize,
113        pc: u32,
114        inst: &Instruction<F>,
115        data: &mut [u8],
116    ) -> Result<Handler<F, Ctx>, StaticProgramError>
117    where
118        Ctx: MeteredExecutionCtxTrait,
119    {
120        let data: &mut E2PreCompute<ShiftPreCompute> = data.borrow_mut();
121        data.chip_idx = chip_idx as u32;
122        let local_opcode = self.pre_compute_impl(pc, inst, &mut data.data)?;
123        dispatch!(execute_e2_handler, local_opcode)
124    }
125}
126
127#[cfg(feature = "aot")]
128impl<F: PrimeField32> AotMeteredExecutor<F> for Rv32Shift256Executor {}
129
130#[inline(always)]
131unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, OP: ShiftOp>(
132    pre_compute: &ShiftPreCompute,
133    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
134) {
135    let rs1_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.b as u32);
136    let rs2_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.c as u32);
137    let rd_ptr = exec_state.vm_read::<u8, 4>(RV32_REGISTER_AS, pre_compute.a as u32);
138    let rs1 = read_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rs1_ptr));
139    let rs2 = read_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rs2_ptr));
140    let rd = OP::compute(rs1, rs2);
141    write_int256(exec_state, RV32_MEMORY_AS, u32::from_le_bytes(rd_ptr), &rd);
142    let pc = exec_state.pc();
143    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
144}
145
146#[create_handler]
147#[inline(always)]
148unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait, OP: ShiftOp>(
149    pre_compute: *const u8,
150    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
151) {
152    let pre_compute: &ShiftPreCompute =
153        std::slice::from_raw_parts(pre_compute, size_of::<ShiftPreCompute>()).borrow();
154    execute_e12_impl::<F, CTX, OP>(pre_compute, exec_state);
155}
156
157#[create_handler]
158#[inline(always)]
159unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait, OP: ShiftOp>(
160    pre_compute: *const u8,
161    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
162) {
163    let pre_compute: &E2PreCompute<ShiftPreCompute> =
164        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<ShiftPreCompute>>())
165            .borrow();
166    exec_state
167        .ctx
168        .on_height_change(pre_compute.chip_idx as usize, 1);
169    execute_e12_impl::<F, CTX, OP>(&pre_compute.data, exec_state);
170}
171
172impl Rv32Shift256Executor {
173    fn pre_compute_impl<F: PrimeField32>(
174        &self,
175        pc: u32,
176        inst: &Instruction<F>,
177        data: &mut ShiftPreCompute,
178    ) -> Result<ShiftOpcode, StaticProgramError> {
179        let Instruction {
180            opcode,
181            a,
182            b,
183            c,
184            d,
185            e,
186            ..
187        } = inst;
188        let e_u32 = e.as_canonical_u32();
189        if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 != RV32_MEMORY_AS {
190            return Err(StaticProgramError::InvalidInstruction(pc));
191        }
192        *data = ShiftPreCompute {
193            a: a.as_canonical_u32() as u8,
194            b: b.as_canonical_u32() as u8,
195            c: c.as_canonical_u32() as u8,
196        };
197        let local_opcode =
198            ShiftOpcode::from_usize(opcode.local_opcode_idx(Rv32Shift256Opcode::CLASS_OFFSET));
199        Ok(local_opcode)
200    }
201}
202
203trait ShiftOp {
204    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> [u8; INT256_NUM_LIMBS];
205}
206struct SllOp;
207struct SrlOp;
208struct SraOp;
209impl ShiftOp for SllOp {
210    #[inline(always)]
211    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> [u8; INT256_NUM_LIMBS] {
212        let rs1_u64: [u64; 4] = bytes_to_u64_array(rs1);
213        let rs2_u64: [u64; 4] = bytes_to_u64_array(rs2);
214        let mut rd = [0u64; 4];
215        // Only use the first 8 bits.
216        let shift = (rs2_u64[0] & 0xff) as u32;
217        let index_offset = (shift / u64::BITS) as usize;
218        let bit_offset = shift % u64::BITS;
219        let mut carry = 0u64;
220        for i in index_offset..4 {
221            let curr = rs1_u64[i - index_offset];
222            rd[i] = (curr << bit_offset) + carry;
223            if bit_offset > 0 {
224                carry = curr >> (u64::BITS - bit_offset);
225            }
226        }
227        u64_array_to_bytes(rd)
228    }
229}
230impl ShiftOp for SrlOp {
231    #[inline(always)]
232    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> [u8; INT256_NUM_LIMBS] {
233        // Logical right shift - fill with 0
234        shift_right(rs1, rs2, 0)
235    }
236}
237impl ShiftOp for SraOp {
238    #[inline(always)]
239    fn compute(rs1: [u8; INT256_NUM_LIMBS], rs2: [u8; INT256_NUM_LIMBS]) -> [u8; INT256_NUM_LIMBS] {
240        // Arithmetic right shift - fill with sign bit
241        if rs1[INT256_NUM_LIMBS - 1] & 0x80 > 0 {
242            shift_right(rs1, rs2, u64::MAX)
243        } else {
244            shift_right(rs1, rs2, 0)
245        }
246    }
247}
248
249#[inline(always)]
250fn shift_right(
251    rs1: [u8; INT256_NUM_LIMBS],
252    rs2: [u8; INT256_NUM_LIMBS],
253    init_value: u64,
254) -> [u8; INT256_NUM_LIMBS] {
255    let rs1_u64: [u64; 4] = bytes_to_u64_array(rs1);
256    let rs2_u64: [u64; 4] = bytes_to_u64_array(rs2);
257    let mut rd = [init_value; 4];
258    let shift = (rs2_u64[0] & 0xff) as u32;
259    let index_offset = (shift / u64::BITS) as usize;
260    let bit_offset = shift % u64::BITS;
261    let mut carry = if bit_offset > 0 {
262        init_value << (u64::BITS - bit_offset)
263    } else {
264        0
265    };
266    for i in (index_offset..4).rev() {
267        let curr = rs1_u64[i];
268        rd[i - index_offset] = (curr >> bit_offset) + carry;
269        if bit_offset > 0 {
270            carry = curr << (u64::BITS - bit_offset);
271        }
272    }
273    u64_array_to_bytes(rd)
274}
275
276#[cfg(test)]
277mod tests {
278    use alloy_primitives::U256;
279    use rand::{prelude::StdRng, Rng, SeedableRng};
280
281    use crate::{
282        shift::{ShiftOp, SllOp, SraOp, SrlOp},
283        INT256_NUM_LIMBS,
284    };
285
286    #[test]
287    fn test_shift_op() {
288        let mut rng = StdRng::from_seed([42; 32]);
289        for _ in 0..10000 {
290            let limbs_a: [u8; INT256_NUM_LIMBS] = rng.random();
291            let mut limbs_b: [u8; INT256_NUM_LIMBS] = [0; INT256_NUM_LIMBS];
292            let shift: u8 = rng.random();
293            limbs_b[0] = shift;
294            let a = U256::from_le_bytes(limbs_a);
295            {
296                let res = SllOp::compute(limbs_a, limbs_b);
297                assert_eq!(U256::from_le_bytes(res), a << shift);
298            }
299            {
300                let res = SraOp::compute(limbs_a, limbs_b);
301                assert_eq!(U256::from_le_bytes(res), a.arithmetic_shr(shift as usize));
302            }
303            {
304                let res = SrlOp::compute(limbs_a, limbs_b);
305                assert_eq!(U256::from_le_bytes(res), a >> shift);
306            }
307        }
308    }
309}
openvm_bigint_circuit/shift.rs

openvm_bigint_circuit/
shift.rs
