1use core::mem::size_of;
23use openvm_circuit::system::memory::offline_checker::{MemoryReadAuxCols, MemoryWriteAuxCols};
4use openvm_circuit_primitives::utils::assert_array_eq;
5use openvm_circuit_primitives_derive::AlignedBorrow;
6use openvm_instructions::riscv::RV32_REGISTER_NUM_LIMBS;
7use openvm_stark_backend::p3_air::AirBuilder;
8use p3_keccak_air::KeccakCols as KeccakPermCols;
910use super::{
11 KECCAK_ABSORB_READS, KECCAK_DIGEST_WRITES, KECCAK_RATE_BYTES, KECCAK_RATE_U16S,
12 KECCAK_REGISTER_READS, KECCAK_WORD_SIZE,
13};
1415#[repr(C)]
16#[derive(Debug, AlignedBorrow)]
17pub struct KeccakVmCols<T> {
18/// Columns for keccak-f permutation
19pub inner: KeccakPermCols<T>,
20/// Columns for sponge and padding
21pub sponge: KeccakSpongeCols<T>,
22/// Columns for instruction interface and register access
23pub instruction: KeccakInstructionCols<T>,
24/// Auxiliary columns for offline memory checking
25pub mem_oc: KeccakMemoryCols<T>,
26}
2728/// Columns for KECCAK256_RV32 instruction parsing.
29/// Includes columns for instruction execution and register reads.
30#[allow(clippy::too_many_arguments)]
31#[repr(C)]
32#[derive(Copy, Clone, Debug, Default, AlignedBorrow, derive_new::new)]
33pub struct KeccakInstructionCols<T> {
34/// Program counter
35pub pc: T,
36/// True for all rows that are part of opcode execution.
37 /// False on dummy rows only used to pad the height.
38pub is_enabled: T,
39/// Is enabled and first round of block. Used to lower constraint degree.
40 /// is_enabled * inner.step_flags\[0\]
41pub is_enabled_first_round: T,
42/// The starting timestamp to use for memory access in this row.
43 /// A single row will do multiple memory accesses.
44pub start_timestamp: T,
45/// Pointer to address space 1 `dst` register
46pub dst_ptr: T,
47/// Pointer to address space 1 `src` register
48pub src_ptr: T,
49/// Pointer to address space 1 `len` register
50pub len_ptr: T,
51// Register values
52/// dst <- \[dst_ptr:4\]_1
53pub dst: [T; RV32_REGISTER_NUM_LIMBS],
54/// src <- \[src_ptr:4\]_1
55 /// We store src_limbs\[i\] = \[src_ptr + i + 1\]_1 and src = u32(\[src_ptr:4\]_1) from which \[src_ptr\]_1
56 /// can be recovered by linear combination.
57 /// We do this because `src` needs to be incremented between keccak-f permutations.
58pub src_limbs: [T; RV32_REGISTER_NUM_LIMBS - 1],
59pub src: T,
60/// len <- \[len_ptr:4\]_1
61 /// We store len_limbs\[i\] = \[len_ptr + i + 1\]_1 and remaining_len = u32(\[len_ptr:4\]_1)
62 /// from which \[len_ptr\]_1 can be recovered by linear combination.
63 /// We do this because `remaining_len` needs to be decremented between keccak-f permutations.
64pub len_limbs: [T; RV32_REGISTER_NUM_LIMBS - 1],
65/// The remaining length of the unpadded input, in bytes.
66 /// If `is_new_start` is true and `is_enabled` is true, this must be equal to `u32(len)`.
67pub remaining_len: T,
68}
6970#[repr(C)]
71#[derive(Clone, Copy, Debug, AlignedBorrow)]
72pub struct KeccakSpongeCols<T> {
73/// Only used on first row of a round to determine whether the state
74 /// prior to absorb should be reset to all 0s.
75 /// Constrained to be zero if not first round.
76pub is_new_start: T,
7778/// Whether the current byte is a padding byte.
79 ///
80 /// If this row represents a full input block, this should contain all 0s.
81pub is_padding_byte: [T; KECCAK_RATE_BYTES],
8283/// The block being absorbed, which may contain input bytes and padding
84 /// bytes.
85pub block_bytes: [T; KECCAK_RATE_BYTES],
8687/// For each of the first [KECCAK_RATE_U16S] `u16` limbs in the state,
88 /// the most significant byte of the limb.
89 /// Here `state` is the postimage state if last round and the preimage
90 /// state if first round. It can be junk if not first or last round.
91pub state_hi: [T; KECCAK_RATE_U16S],
92}
9394#[repr(C)]
95#[derive(Clone, Debug, AlignedBorrow)]
96pub struct KeccakMemoryCols<T> {
97pub register_aux: [MemoryReadAuxCols<T>; KECCAK_REGISTER_READS],
98pub absorb_reads: [MemoryReadAuxCols<T>; KECCAK_ABSORB_READS],
99pub digest_writes: [MemoryWriteAuxCols<T, KECCAK_WORD_SIZE>; KECCAK_DIGEST_WRITES],
100/// The input bytes are batch read in blocks of private constant KECCAK_WORD_SIZE bytes. However
101 /// if the input length is not a multiple of KECCAK_WORD_SIZE, we read into
102 /// `partial_block` more bytes than we need. On the other hand `block_bytes` expects
103 /// only the partial block of bytes and then the correctly padded bytes.
104 /// We will select between `partial_block` and `block_bytes` for what to read from memory.
105 /// We never read a full padding block, so the first byte is always ok.
106pub partial_block: [T; KECCAK_WORD_SIZE - 1],
107}
108109impl<T: Copy> KeccakVmCols<T> {
110pub const fn remaining_len(&self) -> T {
111self.instruction.remaining_len
112 }
113114pub const fn is_new_start(&self) -> T {
115self.sponge.is_new_start
116 }
117118pub fn postimage(&self, y: usize, x: usize, limb: usize) -> T {
119self.inner.a_prime_prime_prime(y, x, limb)
120 }
121122pub fn is_first_round(&self) -> T {
123*self.inner.step_flags.first().unwrap()
124 }
125126pub fn is_last_round(&self) -> T {
127*self.inner.step_flags.last().unwrap()
128 }
129}
130131impl<T: Copy> KeccakInstructionCols<T> {
132pub fn assert_eq<AB: AirBuilder>(&self, builder: &mut AB, other: Self)
133where
134T: Into<AB::Expr>,
135 {
136 builder.assert_eq(self.pc, other.pc);
137 builder.assert_eq(self.is_enabled, other.is_enabled);
138 builder.assert_eq(self.start_timestamp, other.start_timestamp);
139 builder.assert_eq(self.dst_ptr, other.dst_ptr);
140 builder.assert_eq(self.src_ptr, other.src_ptr);
141 builder.assert_eq(self.len_ptr, other.len_ptr);
142 assert_array_eq(builder, self.dst, other.dst);
143 assert_array_eq(builder, self.src_limbs, other.src_limbs);
144 builder.assert_eq(self.src, other.src);
145 assert_array_eq(builder, self.len_limbs, other.len_limbs);
146 builder.assert_eq(self.remaining_len, other.remaining_len);
147 }
148}
149150pub const NUM_KECCAK_VM_COLS: usize = size_of::<KeccakVmCols<u8>>();
151pub const NUM_KECCAK_INSTRUCTION_COLS: usize = size_of::<KeccakInstructionCols<u8>>();
152pub const NUM_KECCAK_SPONGE_COLS: usize = size_of::<KeccakSpongeCols<u8>>();
153pub const NUM_KECCAK_MEMORY_COLS: usize = size_of::<KeccakMemoryCols<u8>>();