openvm_keccak256_circuit/xorin/
execution.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    convert::TryInto,
4    mem::size_of,
5};
6
7use openvm_circuit::{
8    arch::{StaticProgramError, *},
9    system::memory::online::GuestMemory,
10};
11use openvm_circuit_primitives_derive::AlignedBytesBorrow;
12use openvm_instructions::{
13    instruction::Instruction,
14    program::DEFAULT_PC_STEP,
15    riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
16};
17use openvm_stark_backend::p3_field::PrimeField32;
18
19use super::XorinVmExecutor;
20use crate::KECCAK_WORD_SIZE;
21
22#[derive(AlignedBytesBorrow, Clone)]
23#[repr(C)]
24struct XorinPreCompute {
25    a: u8,
26    b: u8,
27    c: u8,
28}
29
30impl XorinVmExecutor {
31    fn pre_compute_impl<F: PrimeField32>(
32        &self,
33        pc: u32,
34        inst: &Instruction<F>,
35        data: &mut XorinPreCompute,
36    ) -> Result<(), StaticProgramError> {
37        let Instruction {
38            opcode: _,
39            a,
40            b,
41            c,
42            d,
43            e,
44            ..
45        } = inst;
46
47        let e_u32 = e.as_canonical_u32();
48        if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 != RV32_MEMORY_AS {
49            return Err(StaticProgramError::InvalidInstruction(pc));
50        }
51
52        *data = XorinPreCompute {
53            a: a.as_canonical_u32() as u8,
54            b: b.as_canonical_u32() as u8,
55            c: c.as_canonical_u32() as u8,
56        };
57
58        Ok(())
59    }
60}
61
62impl<F: PrimeField32> InterpreterExecutor<F> for XorinVmExecutor {
63    fn pre_compute_size(&self) -> usize {
64        size_of::<XorinPreCompute>()
65    }
66
67    #[cfg(not(feature = "tco"))]
68    fn pre_compute<Ctx>(
69        &self,
70        pc: u32,
71        inst: &Instruction<F>,
72        data: &mut [u8],
73    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
74    where
75        Ctx: ExecutionCtxTrait,
76    {
77        let data: &mut XorinPreCompute = data.borrow_mut();
78        self.pre_compute_impl(pc, inst, data)?;
79        Ok(execute_e1_impl::<_, _>)
80    }
81
82    #[cfg(feature = "tco")]
83    fn handler<Ctx>(
84        &self,
85        pc: u32,
86        inst: &Instruction<F>,
87        data: &mut [u8],
88    ) -> Result<Handler<F, Ctx>, StaticProgramError>
89    where
90        Ctx: ExecutionCtxTrait,
91    {
92        let data: &mut XorinPreCompute = data.borrow_mut();
93        self.pre_compute_impl(pc, inst, data)?;
94        Ok(execute_e1_handler)
95    }
96}
97
98#[cfg(feature = "aot")]
99impl<F: PrimeField32> AotExecutor<F> for XorinVmExecutor {}
100
101impl<F: PrimeField32> InterpreterMeteredExecutor<F> for XorinVmExecutor {
102    fn metered_pre_compute_size(&self) -> usize {
103        size_of::<E2PreCompute<XorinPreCompute>>()
104    }
105
106    #[cfg(not(feature = "tco"))]
107    fn metered_pre_compute<Ctx>(
108        &self,
109        chip_idx: usize,
110        pc: u32,
111        inst: &Instruction<F>,
112        data: &mut [u8],
113    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
114    where
115        Ctx: MeteredExecutionCtxTrait,
116    {
117        let data: &mut E2PreCompute<XorinPreCompute> = data.borrow_mut();
118        data.chip_idx = chip_idx as u32;
119        self.pre_compute_impl(pc, inst, &mut data.data)?;
120        Ok(execute_e2_impl::<_, _>)
121    }
122
123    #[cfg(feature = "tco")]
124    fn metered_handler<Ctx>(
125        &self,
126        chip_idx: usize,
127        pc: u32,
128        inst: &Instruction<F>,
129        data: &mut [u8],
130    ) -> Result<Handler<F, Ctx>, StaticProgramError>
131    where
132        Ctx: MeteredExecutionCtxTrait,
133    {
134        let data: &mut E2PreCompute<XorinPreCompute> = data.borrow_mut();
135        data.chip_idx = chip_idx as u32;
136        self.pre_compute_impl(pc, inst, &mut data.data)?;
137        Ok(execute_e2_handler)
138    }
139}
140
141#[cfg(feature = "aot")]
142impl<F: PrimeField32> AotMeteredExecutor<F> for XorinVmExecutor {}
143
144#[create_handler]
145#[inline(always)]
146unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
147    pre_compute: *const u8,
148    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
149) {
150    let pre_compute: &XorinPreCompute =
151        std::slice::from_raw_parts(pre_compute, size_of::<XorinPreCompute>()).borrow();
152    execute_e12_impl::<F, CTX, true>(pre_compute, exec_state);
153}
154
155#[inline(always)]
156unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const IS_E1: bool>(
157    pre_compute: &XorinPreCompute,
158    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
159) {
160    let buffer = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.a as u32);
161    let input = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.b as u32);
162    let length = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.c as u32);
163    let buffer_u32 = u32::from_le_bytes(buffer);
164    let input_u32 = u32::from_le_bytes(input);
165    let length_u32 = u32::from_le_bytes(length);
166
167    // SAFETY: RV32_MEMORY_AS is memory address space of type u8
168    let num_reads = (length_u32 as usize).div_ceil(KECCAK_WORD_SIZE);
169    let buffer_bytes: Vec<_> = (0..num_reads)
170        .flat_map(|i| {
171            exec_state.vm_read::<u8, KECCAK_WORD_SIZE>(
172                RV32_MEMORY_AS,
173                buffer_u32 + (i * KECCAK_WORD_SIZE) as u32,
174            )
175        })
176        .collect();
177
178    let input_bytes: Vec<_> = (0..num_reads)
179        .flat_map(|i| {
180            exec_state.vm_read::<u8, KECCAK_WORD_SIZE>(
181                RV32_MEMORY_AS,
182                input_u32 + (i * KECCAK_WORD_SIZE) as u32,
183            )
184        })
185        .collect();
186
187    let mut output_bytes = buffer_bytes;
188    for i in 0..output_bytes.len() {
189        output_bytes[i] ^= input_bytes[i];
190    }
191
192    // Write XOR result back to the buffer memory in KECCAK_WORD_SIZE chunks.
193    // Note: this means output_bytes has to be multiple of KECCAK_WORD_SIZE
194    // Todo: recheck the above condition is okay
195    for (i, chunk) in output_bytes.chunks_exact(KECCAK_WORD_SIZE).enumerate() {
196        let chunk: [u8; KECCAK_WORD_SIZE] = chunk.try_into().unwrap();
197        exec_state.vm_write::<u8, KECCAK_WORD_SIZE>(
198            RV32_MEMORY_AS,
199            buffer_u32 + (i * KECCAK_WORD_SIZE) as u32,
200            &chunk,
201        );
202    }
203
204    let pc = exec_state.pc();
205    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
206}
207
208#[create_handler]
209#[inline(always)]
210unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait>(
211    pre_compute: *const u8,
212    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
213) {
214    let pre_compute: &E2PreCompute<XorinPreCompute> =
215        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<XorinPreCompute>>())
216            .borrow();
217    execute_e12_impl::<F, CTX, false>(&pre_compute.data, exec_state);
218    exec_state
219        .ctx
220        .on_height_change(pre_compute.chip_idx as usize, 1);
221}
openvm_keccak256_circuit/xorin/execution.rs

openvm_keccak256_circuit/xorin/
execution.rs
