openvm_keccak256_circuit/
lib.rs

1//! Stateful keccak256 hasher. Handles full keccak sponge (padding, absorb, keccak-f) on
2//! variable length inputs read from VM memory.
3use std::{
4    array::from_fn,
5    cmp::min,
6    sync::{Arc, Mutex},
7};
8
9use openvm_circuit_primitives::bitwise_op_lookup::SharedBitwiseOperationLookupChip;
10use openvm_stark_backend::p3_field::PrimeField32;
11use serde::{Deserialize, Serialize};
12use serde_big_array::BigArray;
13use tiny_keccak::{Hasher, Keccak};
14use utils::num_keccak_f;
15
16pub mod air;
17pub mod columns;
18pub mod trace;
19pub mod utils;
20
21mod extension;
22pub use extension::*;
23
24#[cfg(test)]
25mod tests;
26
27pub use air::KeccakVmAir;
28use openvm_circuit::{
29    arch::{ExecutionBridge, ExecutionBus, ExecutionError, ExecutionState, InstructionExecutor},
30    system::{
31        memory::{offline_checker::MemoryBridge, MemoryController, OfflineMemory, RecordId},
32        program::ProgramBus,
33    },
34};
35use openvm_instructions::{
36    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_NUM_LIMBS, LocalOpcode,
37};
38use openvm_keccak256_transpiler::Rv32KeccakOpcode;
39use openvm_rv32im_circuit::adapters::read_rv32_register;
40
41// ==== Constants for register/memory adapter ====
42/// Register reads to get dst, src, len
43const KECCAK_REGISTER_READS: usize = 3;
44/// Number of cells to read/write in a single memory access
45const KECCAK_WORD_SIZE: usize = 4;
46/// Memory reads for absorb per row
47const KECCAK_ABSORB_READS: usize = KECCAK_RATE_BYTES / KECCAK_WORD_SIZE;
48/// Memory writes for digest per row
49const KECCAK_DIGEST_WRITES: usize = KECCAK_DIGEST_BYTES / KECCAK_WORD_SIZE;
50
51// ==== Do not change these constants! ====
52/// Total number of sponge bytes: number of rate bytes + number of capacity
53/// bytes.
54pub const KECCAK_WIDTH_BYTES: usize = 200;
55/// Total number of 16-bit limbs in the sponge.
56pub const KECCAK_WIDTH_U16S: usize = KECCAK_WIDTH_BYTES / 2;
57/// Number of rate bytes.
58pub const KECCAK_RATE_BYTES: usize = 136;
59/// Number of 16-bit rate limbs.
60pub const KECCAK_RATE_U16S: usize = KECCAK_RATE_BYTES / 2;
61/// Number of absorb rounds, equal to rate in u64s.
62pub const NUM_ABSORB_ROUNDS: usize = KECCAK_RATE_BYTES / 8;
63/// Number of capacity bytes.
64pub const KECCAK_CAPACITY_BYTES: usize = 64;
65/// Number of 16-bit capacity limbs.
66pub const KECCAK_CAPACITY_U16S: usize = KECCAK_CAPACITY_BYTES / 2;
67/// Number of output digest bytes used during the squeezing phase.
68pub const KECCAK_DIGEST_BYTES: usize = 32;
69/// Number of 64-bit digest limbs.
70pub const KECCAK_DIGEST_U64S: usize = KECCAK_DIGEST_BYTES / 8;
71
72pub struct KeccakVmChip<F: PrimeField32> {
73    pub air: KeccakVmAir,
74    /// IO and memory data necessary for each opcode call
75    pub records: Vec<KeccakRecord<F>>,
76    pub bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
77
78    offset: usize,
79
80    offline_memory: Arc<Mutex<OfflineMemory<F>>>,
81}
82
83#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
84pub struct KeccakRecord<F> {
85    pub pc: F,
86    pub dst_read: RecordId,
87    pub src_read: RecordId,
88    pub len_read: RecordId,
89    pub input_blocks: Vec<KeccakInputBlock>,
90    pub digest_writes: [RecordId; KECCAK_DIGEST_WRITES],
91}
92
93#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
94pub struct KeccakInputBlock {
95    /// Memory reads for non-padding bytes in this block. Length is at most [KECCAK_RATE_BYTES /
96    /// KECCAK_WORD_SIZE].
97    pub reads: Vec<RecordId>,
98    /// Index in `reads` of the memory read for < KECCAK_WORD_SIZE bytes, if any.
99    pub partial_read_idx: Option<usize>,
100    /// Bytes with padding. Can be derived from `bytes_read` but we store for convenience.
101    #[serde(with = "BigArray")]
102    pub padded_bytes: [u8; KECCAK_RATE_BYTES],
103    pub remaining_len: usize,
104    pub src: usize,
105    pub is_new_start: bool,
106}
107
108impl<F: PrimeField32> KeccakVmChip<F> {
109    pub fn new(
110        execution_bus: ExecutionBus,
111        program_bus: ProgramBus,
112        memory_bridge: MemoryBridge,
113        address_bits: usize,
114        bitwise_lookup_chip: SharedBitwiseOperationLookupChip<8>,
115        offset: usize,
116        offline_memory: Arc<Mutex<OfflineMemory<F>>>,
117    ) -> Self {
118        Self {
119            air: KeccakVmAir::new(
120                ExecutionBridge::new(execution_bus, program_bus),
121                memory_bridge,
122                bitwise_lookup_chip.bus(),
123                address_bits,
124                offset,
125            ),
126            bitwise_lookup_chip,
127            records: Vec::new(),
128            offset,
129            offline_memory,
130        }
131    }
132}
133
134impl<F: PrimeField32> InstructionExecutor<F> for KeccakVmChip<F> {
135    fn execute(
136        &mut self,
137        memory: &mut MemoryController<F>,
138        instruction: &Instruction<F>,
139        from_state: ExecutionState<u32>,
140    ) -> Result<ExecutionState<u32>, ExecutionError> {
141        let &Instruction {
142            opcode,
143            a,
144            b,
145            c,
146            d,
147            e,
148            ..
149        } = instruction;
150        let local_opcode = Rv32KeccakOpcode::from_usize(opcode.local_opcode_idx(self.offset));
151        debug_assert_eq!(local_opcode, Rv32KeccakOpcode::KECCAK256);
152
153        let mut timestamp_delta = 3;
154        let (dst_read, dst) = read_rv32_register(memory, d, a);
155        let (src_read, src) = read_rv32_register(memory, d, b);
156        let (len_read, len) = read_rv32_register(memory, d, c);
157        #[cfg(debug_assertions)]
158        {
159            assert!(dst < (1 << self.air.ptr_max_bits));
160            assert!(src < (1 << self.air.ptr_max_bits));
161            assert!(len < (1 << self.air.ptr_max_bits));
162        }
163
164        let mut remaining_len = len as usize;
165        let num_blocks = num_keccak_f(remaining_len);
166        let mut input_blocks = Vec::with_capacity(num_blocks);
167        let mut hasher = Keccak::v256();
168        let mut src = src as usize;
169
170        for block_idx in 0..num_blocks {
171            if block_idx != 0 {
172                memory.increment_timestamp_by(KECCAK_REGISTER_READS as u32);
173                timestamp_delta += KECCAK_REGISTER_READS as u32;
174            }
175            let mut reads = Vec::with_capacity(KECCAK_RATE_BYTES);
176
177            let mut partial_read_idx = None;
178            let mut bytes = [0u8; KECCAK_RATE_BYTES];
179            for i in (0..KECCAK_RATE_BYTES).step_by(KECCAK_WORD_SIZE) {
180                if i < remaining_len {
181                    let read =
182                        memory.read::<RV32_REGISTER_NUM_LIMBS>(e, F::from_canonical_usize(src + i));
183
184                    let chunk = read.1.map(|x| {
185                        x.as_canonical_u32()
186                            .try_into()
187                            .expect("Memory cell not a byte")
188                    });
189                    let copy_len = min(KECCAK_WORD_SIZE, remaining_len - i);
190                    if copy_len != KECCAK_WORD_SIZE {
191                        partial_read_idx = Some(reads.len());
192                    }
193                    bytes[i..i + copy_len].copy_from_slice(&chunk[..copy_len]);
194                    reads.push(read.0);
195                } else {
196                    memory.increment_timestamp();
197                }
198                timestamp_delta += 1;
199            }
200
201            let mut block = KeccakInputBlock {
202                reads,
203                partial_read_idx,
204                padded_bytes: bytes,
205                remaining_len,
206                src,
207                is_new_start: block_idx == 0,
208            };
209            if block_idx != num_blocks - 1 {
210                src += KECCAK_RATE_BYTES;
211                remaining_len -= KECCAK_RATE_BYTES;
212                hasher.update(&block.padded_bytes);
213            } else {
214                // handle padding here since it is convenient
215                debug_assert!(remaining_len < KECCAK_RATE_BYTES);
216                hasher.update(&block.padded_bytes[..remaining_len]);
217
218                if remaining_len == KECCAK_RATE_BYTES - 1 {
219                    block.padded_bytes[remaining_len] = 0b1000_0001;
220                } else {
221                    block.padded_bytes[remaining_len] = 0x01;
222                    block.padded_bytes[KECCAK_RATE_BYTES - 1] = 0x80;
223                }
224            }
225            input_blocks.push(block);
226        }
227        let mut output = [0u8; 32];
228        hasher.finalize(&mut output);
229        let dst = dst as usize;
230        let digest_writes: [_; KECCAK_DIGEST_WRITES] = from_fn(|i| {
231            timestamp_delta += 1;
232            memory
233                .write::<KECCAK_WORD_SIZE>(
234                    e,
235                    F::from_canonical_usize(dst + i * KECCAK_WORD_SIZE),
236                    from_fn(|j| F::from_canonical_u8(output[i * KECCAK_WORD_SIZE + j])),
237                )
238                .0
239        });
240        tracing::trace!("[runtime] keccak256 output: {:?}", output);
241
242        let record = KeccakRecord {
243            pc: F::from_canonical_u32(from_state.pc),
244            dst_read,
245            src_read,
246            len_read,
247            input_blocks,
248            digest_writes,
249        };
250
251        // Add the events to chip state for later trace generation usage
252        self.records.push(record);
253
254        // NOTE: Check this is consistent with KeccakVmAir::timestamp_change (we don't use it to
255        // avoid unnecessary conversions here)
256        let total_timestamp_delta =
257            len + (KECCAK_REGISTER_READS + KECCAK_ABSORB_READS + KECCAK_DIGEST_WRITES) as u32;
258        memory.increment_timestamp_by(total_timestamp_delta - timestamp_delta);
259
260        Ok(ExecutionState {
261            pc: from_state.pc + DEFAULT_PC_STEP,
262            timestamp: from_state.timestamp + total_timestamp_delta,
263        })
264    }
265
266    fn get_opcode_name(&self, _: usize) -> String {
267        "KECCAK256".to_string()
268    }
269}
270
271impl Default for KeccakInputBlock {
272    fn default() -> Self {
273        // Padding for empty byte array so padding constraints still hold
274        let mut padded_bytes = [0u8; KECCAK_RATE_BYTES];
275        padded_bytes[0] = 0x01;
276        *padded_bytes.last_mut().unwrap() = 0x80;
277        Self {
278            padded_bytes,
279            partial_read_idx: None,
280            remaining_len: 0,
281            is_new_start: true,
282            reads: Vec::new(),
283            src: 0,
284        }
285    }
286}
openvm_keccak256_circuit/lib.rs

openvm_keccak256_circuit/
lib.rs
