openvm_sha2_circuit/sha2_chips/
execution.rs

1use std::borrow::{Borrow, BorrowMut};
2
3use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
4use openvm_circuit_primitives::AlignedBytesBorrow;
5use openvm_instructions::{
6    instruction::Instruction,
7    program::DEFAULT_PC_STEP,
8    riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
9    LocalOpcode,
10};
11use openvm_stark_backend::p3_field::PrimeField32;
12
13use super::{Sha2Config, Sha2VmExecutor, SHA2_READ_SIZE};
14use crate::SHA2_WRITE_SIZE;
15
16#[derive(AlignedBytesBorrow, Clone)]
17#[repr(C)]
18struct Sha2PreCompute {
19    a: u8,
20    b: u8,
21    c: u8,
22}
23
24impl<F: PrimeField32, C: Sha2Config> InterpreterExecutor<F> for Sha2VmExecutor<C> {
25    fn pre_compute_size(&self) -> usize {
26        size_of::<Sha2PreCompute>()
27    }
28
29    #[cfg(not(feature = "tco"))]
30    fn pre_compute<Ctx>(
31        &self,
32        pc: u32,
33        inst: &Instruction<F>,
34        data: &mut [u8],
35    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
36    where
37        Ctx: ExecutionCtxTrait,
38    {
39        let data: &mut Sha2PreCompute = data.borrow_mut();
40        self.pre_compute_impl(pc, inst, data)?;
41        Ok(execute_e1_impl::<_, _, C>)
42    }
43
44    #[cfg(feature = "tco")]
45    fn handler<Ctx>(
46        &self,
47        pc: u32,
48        inst: &Instruction<F>,
49        data: &mut [u8],
50    ) -> Result<Handler<F, Ctx>, StaticProgramError>
51    where
52        Ctx: ExecutionCtxTrait,
53    {
54        let data: &mut Sha2PreCompute = data.borrow_mut();
55        self.pre_compute_impl(pc, inst, data)?;
56        Ok(execute_e1_handler::<_, _, C>)
57    }
58}
59
60impl<F: PrimeField32, C: Sha2Config> InterpreterMeteredExecutor<F> for Sha2VmExecutor<C> {
61    fn metered_pre_compute_size(&self) -> usize {
62        size_of::<E2PreCompute<Sha2PreCompute>>()
63    }
64
65    #[cfg(not(feature = "tco"))]
66    fn metered_pre_compute<Ctx>(
67        &self,
68        chip_idx: usize,
69        pc: u32,
70        inst: &Instruction<F>,
71        data: &mut [u8],
72    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
73    where
74        Ctx: MeteredExecutionCtxTrait,
75    {
76        let data: &mut E2PreCompute<Sha2PreCompute> = data.borrow_mut();
77        data.chip_idx = chip_idx as u32;
78        self.pre_compute_impl(pc, inst, &mut data.data)?;
79        Ok(execute_e2_impl::<_, _, C>)
80    }
81
82    #[cfg(feature = "tco")]
83    fn metered_handler<Ctx>(
84        &self,
85        chip_idx: usize,
86        pc: u32,
87        inst: &Instruction<F>,
88        data: &mut [u8],
89    ) -> Result<Handler<F, Ctx>, StaticProgramError>
90    where
91        Ctx: MeteredExecutionCtxTrait,
92    {
93        let data: &mut E2PreCompute<Sha2PreCompute> = data.borrow_mut();
94        data.chip_idx = chip_idx as u32;
95        self.pre_compute_impl(pc, inst, &mut data.data)?;
96        Ok(execute_e2_handler::<_, _, C>)
97    }
98}
99
100#[inline(always)]
101unsafe fn execute_e12_impl<
102    F: PrimeField32,
103    C: Sha2Config,
104    CTX: ExecutionCtxTrait,
105    const IS_E1: bool,
106>(
107    pre_compute: &Sha2PreCompute,
108    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
109) -> u32 {
110    let dst = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.a as u32);
111    let state = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.b as u32);
112    let input = exec_state.vm_read(RV32_REGISTER_AS, pre_compute.c as u32);
113    let dst_u32 = u32::from_le_bytes(dst);
114    let state_u32 = u32::from_le_bytes(state);
115    let input_u32 = u32::from_le_bytes(input);
116
117    // state is in 4-byte little-endian words
118    let mut state_data = Vec::with_capacity(C::STATE_BYTES);
119    for i in 0..C::STATE_READS {
120        state_data.extend_from_slice(&exec_state.vm_read::<u8, SHA2_READ_SIZE>(
121            RV32_MEMORY_AS,
122            state_u32 + (i * SHA2_READ_SIZE) as u32,
123        ));
124    }
125    let mut input_block = Vec::with_capacity(C::BLOCK_BYTES);
126    for i in 0..C::BLOCK_READS {
127        input_block.extend_from_slice(&exec_state.vm_read::<u8, SHA2_READ_SIZE>(
128            RV32_MEMORY_AS,
129            input_u32 + (i * SHA2_READ_SIZE) as u32,
130        ));
131    }
132
133    C::compress(&mut state_data, &input_block);
134
135    for i in 0..C::STATE_WRITES {
136        exec_state.vm_write::<u8, SHA2_WRITE_SIZE>(
137            RV32_MEMORY_AS,
138            dst_u32 + (i * SHA2_WRITE_SIZE) as u32,
139            &state_data[i * SHA2_WRITE_SIZE..(i + 1) * SHA2_WRITE_SIZE]
140                .try_into()
141                .unwrap(),
142        );
143    }
144
145    let pc = exec_state.pc();
146    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
147
148    1 // height delta
149}
150
151#[create_handler]
152#[inline(always)]
153unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait, C: Sha2Config>(
154    pre_compute: *const u8,
155    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
156) {
157    let pre_compute: &Sha2PreCompute =
158        std::slice::from_raw_parts(pre_compute, size_of::<Sha2PreCompute>()).borrow();
159    execute_e12_impl::<F, C, CTX, true>(pre_compute, exec_state);
160}
161
162#[create_handler]
163#[inline(always)]
164unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait, C: Sha2Config>(
165    pre_compute: *const u8,
166    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
167) {
168    let pre_compute: &E2PreCompute<Sha2PreCompute> =
169        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<Sha2PreCompute>>()).borrow();
170
171    let main_air_idx = pre_compute.chip_idx as usize;
172
173    // Update Sha2MainChip height (1 row per instruction)
174    let height = execute_e12_impl::<F, C, CTX, false>(&pre_compute.data, exec_state);
175    exec_state.ctx.on_height_change(main_air_idx, height);
176
177    // HACK: Sha2BlockHasherVmAir is added right before Sha2MainAir in extend_circuit,
178    // and due to reverse ordering of AIR indices, block_hasher_air_idx = main_air_idx + 1.
179    // See extension/mod.rs extend_circuit for the ordering.
180    let block_hasher_air_idx = main_air_idx + 1;
181
182    // Update Sha2BlockHasherChip height (ROWS_PER_BLOCK rows per block compression)
183    exec_state
184        .ctx
185        .on_height_change(block_hasher_air_idx, C::ROWS_PER_BLOCK as u32);
186}
187
188#[cfg(feature = "aot")]
189impl<F: PrimeField32, C: Sha2Config> AotExecutor<F> for Sha2VmExecutor<C> {}
190
191#[cfg(feature = "aot")]
192impl<F: PrimeField32, C: Sha2Config> AotMeteredExecutor<F> for Sha2VmExecutor<C> {}
193
194impl<C: Sha2Config> Sha2VmExecutor<C> {
195    fn pre_compute_impl<F: PrimeField32>(
196        &self,
197        pc: u32,
198        inst: &Instruction<F>,
199        data: &mut Sha2PreCompute,
200    ) -> Result<(), StaticProgramError> {
201        let Instruction {
202            opcode,
203            a,
204            b,
205            c,
206            d,
207            e,
208            ..
209        } = inst;
210        let e_u32 = e.as_canonical_u32();
211        if d.as_canonical_u32() != RV32_REGISTER_AS || e_u32 != RV32_MEMORY_AS {
212            return Err(StaticProgramError::InvalidInstruction(pc));
213        }
214        *data = Sha2PreCompute {
215            a: a.as_canonical_u32() as u8,
216            b: b.as_canonical_u32() as u8,
217            c: c.as_canonical_u32() as u8,
218        };
219        assert_eq!(&C::OPCODE.global_opcode(), opcode);
220        Ok(())
221    }
222}
openvm_sha2_circuit/sha2_chips/execution.rs

openvm_sha2_circuit/sha2_chips/
execution.rs
