openvm_keccak256_circuit/
columns.rs

1use core::mem::size_of;
2
3use 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;
9
10use super::{
11    KECCAK_ABSORB_READS, KECCAK_DIGEST_WRITES, KECCAK_RATE_BYTES, KECCAK_RATE_U16S,
12    KECCAK_REGISTER_READS, KECCAK_WORD_SIZE,
13};
14
15#[repr(C)]
16#[derive(Debug, AlignedBorrow)]
17pub struct KeccakVmCols<T> {
18    /// Columns for keccak-f permutation
19    pub inner: KeccakPermCols<T>,
20    /// Columns for sponge and padding
21    pub sponge: KeccakSpongeCols<T>,
22    /// Columns for instruction interface and register access
23    pub instruction: KeccakInstructionCols<T>,
24    /// Auxiliary columns for offline memory checking
25    pub mem_oc: KeccakMemoryCols<T>,
26}
27
28/// 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
35    pub pc: T,
36    /// True for all rows that are part of opcode execution.
37    /// False on dummy rows only used to pad the height.
38    pub is_enabled: T,
39    /// Is enabled and first round of block. Used to lower constraint degree.
40    /// is_enabled * inner.step_flags\[0\]
41    pub 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.
44    pub start_timestamp: T,
45    /// Pointer to address space 1 `dst` register
46    pub dst_ptr: T,
47    /// Pointer to address space 1 `src` register
48    pub src_ptr: T,
49    /// Pointer to address space 1 `len` register
50    pub len_ptr: T,
51    // Register values
52    /// dst <- \[dst_ptr:4\]_1
53    pub 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.
58    pub src_limbs: [T; RV32_REGISTER_NUM_LIMBS - 1],
59    pub 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.
64    pub 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)`.
67    pub remaining_len: T,
68}
69
70#[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.
76    pub is_new_start: T,
77
78    /// Whether the current byte is a padding byte.
79    ///
80    /// If this row represents a full input block, this should contain all 0s.
81    pub is_padding_byte: [T; KECCAK_RATE_BYTES],
82
83    /// The block being absorbed, which may contain input bytes and padding
84    /// bytes.
85    pub block_bytes: [T; KECCAK_RATE_BYTES],
86
87    /// 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.
91    pub state_hi: [T; KECCAK_RATE_U16S],
92}
93
94#[repr(C)]
95#[derive(Clone, Debug, AlignedBorrow)]
96pub struct KeccakMemoryCols<T> {
97    pub register_aux: [MemoryReadAuxCols<T>; KECCAK_REGISTER_READS],
98    pub absorb_reads: [MemoryReadAuxCols<T>; KECCAK_ABSORB_READS],
99    pub 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.
106    pub partial_block: [T; KECCAK_WORD_SIZE - 1],
107}
108
109impl<T: Copy> KeccakVmCols<T> {
110    pub const fn remaining_len(&self) -> T {
111        self.instruction.remaining_len
112    }
113
114    pub const fn is_new_start(&self) -> T {
115        self.sponge.is_new_start
116    }
117
118    pub fn postimage(&self, y: usize, x: usize, limb: usize) -> T {
119        self.inner.a_prime_prime_prime(y, x, limb)
120    }
121
122    pub fn is_first_round(&self) -> T {
123        *self.inner.step_flags.first().unwrap()
124    }
125
126    pub fn is_last_round(&self) -> T {
127        *self.inner.step_flags.last().unwrap()
128    }
129}
130
131impl<T: Copy> KeccakInstructionCols<T> {
132    pub fn assert_eq<AB: AirBuilder>(&self, builder: &mut AB, other: Self)
133    where
134        T: 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}
149
150pub 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>>();