openvm_keccak256_circuit/
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,
11    riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
12    LocalOpcode,
13};
14use openvm_keccak256_transpiler::Rv32KeccakOpcode;
15use openvm_stark_backend::p3_field::PrimeField32;
16use p3_keccak_air::NUM_ROUNDS;
17
18use super::{KeccakVmExecutor, KECCAK_WORD_SIZE};
19use crate::utils::{keccak256, num_keccak_f};
20
21#[derive(AlignedBytesBorrow, Clone)]
22#[repr(C)]
23struct KeccakPreCompute {
24    a: u8,
25    b: u8,
26    c: u8,
27}
28
29impl KeccakVmExecutor {
30    fn pre_compute_impl<F: PrimeField32>(
31        &self,
32        pc: u32,
33        inst: &Instruction<F>,
34        data: &mut KeccakPreCompute,
35    ) -> Result<(), StaticProgramError> {
36        let Instruction {
37            opcode,
38            a,
39            b,
40            c,
41            d,
42            e,
43            ..
44        } = inst;
45        let e_u32 = e.as_canonical_u32();
46        if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 != RV32_MEMORY_AS {
47            return Err(StaticProgramError::InvalidInstruction(pc));
48        }
49        *data = KeccakPreCompute {
50            a: a.as_canonical_u32() as u8,
51            b: b.as_canonical_u32() as u8,
52            c: c.as_canonical_u32() as u8,
53        };
54        assert_eq!(&Rv32KeccakOpcode::KECCAK256.global_opcode(), opcode);
55        Ok(())
56    }
57}
58
59impl<F: PrimeField32> Executor<F> for KeccakVmExecutor {
60    fn pre_compute_size(&self) -> usize {
61        size_of::<KeccakPreCompute>()
62    }
63
64    #[cfg(not(feature = "tco"))]
65    fn pre_compute<Ctx>(
66        &self,
67        pc: u32,
68        inst: &Instruction<F>,
69        data: &mut [u8],
70    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
71    where
72        Ctx: ExecutionCtxTrait,
73    {
74        let data: &mut KeccakPreCompute = data.borrow_mut();
75        self.pre_compute_impl(pc, inst, data)?;
76        Ok(execute_e1_impl::<_, _>)
77    }
78
79    #[cfg(feature = "tco")]
80    fn handler<Ctx>(
81        &self,
82        pc: u32,
83        inst: &Instruction<F>,
84        data: &mut [u8],
85    ) -> Result<Handler<F, Ctx>, StaticProgramError>
86    where
87        Ctx: ExecutionCtxTrait,
88    {
89        let data: &mut KeccakPreCompute = data.borrow_mut();
90        self.pre_compute_impl(pc, inst, data)?;
91        Ok(execute_e1_handler)
92    }
93}
94
95impl<F: PrimeField32> MeteredExecutor<F> for KeccakVmExecutor {
96    fn metered_pre_compute_size(&self) -> usize {
97        size_of::<E2PreCompute<KeccakPreCompute>>()
98    }
99
100    #[cfg(not(feature = "tco"))]
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<KeccakPreCompute> = data.borrow_mut();
112        data.chip_idx = chip_idx as u32;
113        self.pre_compute_impl(pc, inst, &mut data.data)?;
114        Ok(execute_e2_impl::<_, _>)
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<KeccakPreCompute> = data.borrow_mut();
129        data.chip_idx = chip_idx as u32;
130        self.pre_compute_impl(pc, inst, &mut data.data)?;
131        Ok(execute_e2_handler::<_, _>)
132    }
133}
134
135#[inline(always)]
136unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait, const IS_E1: bool>(
137    pre_compute: &KeccakPreCompute,
138    instret: &mut u64,
139    pc: &mut u32,
140    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
141) -> u32 {
142    let dst = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.a as u32);
143    let src = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.b as u32);
144    let len = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.c as u32);
145    let dst_u32 = u32::from_le_bytes(dst);
146    let src_u32 = u32::from_le_bytes(src);
147    let len_u32 = u32::from_le_bytes(len);
148
149    let (output, height) = if IS_E1 {
150        // SAFETY: RV32_MEMORY_AS is memory address space of type u8
151        let message = exec_state.vm_read_slice(RV32_MEMORY_AS, src_u32, len_u32 as usize);
152        let output = keccak256(message);
153        (output, 0)
154    } else {
155        let num_reads = (len_u32 as usize).div_ceil(KECCAK_WORD_SIZE);
156        let message: Vec<_> = (0..num_reads)
157            .flat_map(|i| {
158                exec_state.vm_read::<u8, KECCAK_WORD_SIZE>(
159                    RV32_MEMORY_AS,
160                    src_u32 + (i * KECCAK_WORD_SIZE) as u32,
161                )
162            })
163            .collect();
164        let output = keccak256(&message[..len_u32 as usize]);
165        let height = (num_keccak_f(len_u32 as usize) * NUM_ROUNDS) as u32;
166        (output, height)
167    };
168    exec_state.vm_write(RV32_MEMORY_AS, dst_u32, &output);
169
170    *pc = pc.wrapping_add(DEFAULT_PC_STEP);
171    *instret += 1;
172
173    height
174}
175
176#[create_handler]
177#[inline(always)]
178unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
179    pre_compute: &[u8],
180    instret: &mut u64,
181    pc: &mut u32,
182    _instret_end: u64,
183    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
184) {
185    let pre_compute: &KeccakPreCompute = pre_compute.borrow();
186    execute_e12_impl::<F, CTX, true>(pre_compute, instret, pc, exec_state);
187}
188
189#[create_handler]
190#[inline(always)]
191unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait>(
192    pre_compute: &[u8],
193    instret: &mut u64,
194    pc: &mut u32,
195    _arg: u64,
196    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
197) {
198    let pre_compute: &E2PreCompute<KeccakPreCompute> = pre_compute.borrow();
199    let height = execute_e12_impl::<F, CTX, false>(&pre_compute.data, instret, pc, exec_state);
200    exec_state
201        .ctx
202        .on_height_change(pre_compute.chip_idx as usize, height);
203}
openvm_keccak256_circuit/execution.rs

openvm_keccak256_circuit/
execution.rs
