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

openvm_sha256_circuit/sha256_chip/
execution.rs
