o1vm/interpreters/keccak/column.rs
1//! This module defines the custom columns used in the Keccak witness, which
2//! are aliases for the actual Keccak witness columns also defined here.
3use self::{Absorbs::*, Sponges::*, Steps::*};
4use crate::{
5 interpreters::keccak::{ZKVM_KECCAK_COLS_CURR, ZKVM_KECCAK_COLS_NEXT},
6 RelationColumnType,
7};
8use kimchi::circuits::polynomials::keccak::constants::{
9 CHI_SHIFTS_B_LEN, CHI_SHIFTS_B_OFF, CHI_SHIFTS_SUM_LEN, CHI_SHIFTS_SUM_OFF, PIRHO_DENSE_E_LEN,
10 PIRHO_DENSE_E_OFF, PIRHO_DENSE_ROT_E_LEN, PIRHO_DENSE_ROT_E_OFF, PIRHO_EXPAND_ROT_E_LEN,
11 PIRHO_EXPAND_ROT_E_OFF, PIRHO_QUOTIENT_E_LEN, PIRHO_QUOTIENT_E_OFF, PIRHO_REMAINDER_E_LEN,
12 PIRHO_REMAINDER_E_OFF, PIRHO_SHIFTS_E_LEN, PIRHO_SHIFTS_E_OFF, QUARTERS, RATE_IN_BYTES,
13 SPONGE_BYTES_LEN, SPONGE_BYTES_OFF, SPONGE_COLS, SPONGE_NEW_STATE_LEN, SPONGE_NEW_STATE_OFF,
14 SPONGE_SHIFTS_LEN, SPONGE_SHIFTS_OFF, SPONGE_ZEROS_LEN, SPONGE_ZEROS_OFF, STATE_LEN,
15 THETA_DENSE_C_LEN, THETA_DENSE_C_OFF, THETA_DENSE_ROT_C_LEN, THETA_DENSE_ROT_C_OFF,
16 THETA_EXPAND_ROT_C_LEN, THETA_EXPAND_ROT_C_OFF, THETA_QUOTIENT_C_LEN, THETA_QUOTIENT_C_OFF,
17 THETA_REMAINDER_C_LEN, THETA_REMAINDER_C_OFF, THETA_SHIFTS_C_LEN, THETA_SHIFTS_C_OFF,
18};
19use kimchi_msm::{
20 columns::{Column, ColumnIndexer},
21 witness::Witness,
22};
23use std::ops::{Index, IndexMut};
24use strum::{EnumCount, IntoEnumIterator};
25use strum_macros::{EnumCount, EnumIter};
26
27/// The maximum total number of witness columns used by the Keccak circuit.
28/// Note that in round steps, the columns used to store padding information are not needed.
29pub const N_ZKVM_KECCAK_REL_COLS: usize = STATUS_LEN + CURR_LEN + NEXT_LEN + RC_LEN;
30
31/// The number of columns required for the Keccak selectors. They are located after the relation columns.
32pub const N_ZKVM_KECCAK_SEL_COLS: usize = 6;
33
34/// Total number of columns used in Keccak, including relation and selectors
35pub const N_ZKVM_KECCAK_COLS: usize = N_ZKVM_KECCAK_REL_COLS + N_ZKVM_KECCAK_SEL_COLS;
36
37const STATUS_OFF: usize = 0; // The offset of the columns reserved for the status indices
38const STATUS_LEN: usize = 3; // The number of columns used by the Keccak circuit to represent the status flags.
39
40const CURR_OFF: usize = STATUS_OFF + STATUS_LEN; // The offset of the curr chunk inside the witness columns
41const CURR_LEN: usize = ZKVM_KECCAK_COLS_CURR; // The length of the curr chunk inside the witness columns
42const NEXT_OFF: usize = CURR_OFF + CURR_LEN; // The offset of the next chunk inside the witness columns
43const NEXT_LEN: usize = ZKVM_KECCAK_COLS_NEXT; // The length of the next chunk inside the witness columns
44
45/// The number of sparse round constants used per round
46pub(crate) const ROUND_CONST_LEN: usize = QUARTERS;
47const RC_OFF: usize = NEXT_OFF + NEXT_LEN; // The offset of the Round coefficients inside the witness columns
48const RC_LEN: usize = ROUND_CONST_LEN + 1; // The round constants plus the round number
49
50const PAD_FLAGS_OFF: usize = STATUS_LEN + SPONGE_COLS; // Offset of the Pad flags inside the witness columns. Starts after sponge columns are finished.
51const PAD_LEN_OFF: usize = 0; // Offset of the PadLength column inside the sponge coefficients
52const PAD_TWO_OFF: usize = 1; // Offset of the TwoToPad column inside the sponge coefficients
53const PAD_SUFFIX_OFF: usize = 2; // Offset of the PadSuffix column inside the sponge coefficients
54/// The padding suffix of 1088 bits is stored as 5 field elements: 1x12 + 4x31 bytes
55pub(crate) const PAD_SUFFIX_LEN: usize = 5;
56const PAD_BYTES_OFF: usize = PAD_SUFFIX_OFF + PAD_SUFFIX_LEN; // Offset of the PadBytesFlags inside the sponge coefficients
57/// The maximum number of padding bytes involved
58pub(crate) const PAD_BYTES_LEN: usize = RATE_IN_BYTES;
59
60const FLAG_ROUND_OFF: usize = N_ZKVM_KECCAK_REL_COLS; // Offset of the Round selector inside DynamicSelector
61const FLAG_FST_OFF: usize = FLAG_ROUND_OFF + 1; // Offset of the Absorb(First) selector inside DynamicSelector
62const FLAG_MID_OFF: usize = FLAG_ROUND_OFF + 2; // Offset of the Absorb(Middle) selector inside DynamicSelector
63const FLAG_LST_OFF: usize = FLAG_ROUND_OFF + 3; // Offset of the Absorb(Last) selector inside DynamicSelector
64const FLAG_ONE_OFF: usize = FLAG_ROUND_OFF + 4; // Offset of the Absorb(Only) selector inside DynamicSelector
65const FLAG_SQUEEZE_OFF: usize = FLAG_ROUND_OFF + 5; // Offset of the Squeeze selector inside DynamicSelector
66
67/// Column aliases used by the Keccak circuit.
68/// The number of aliases is not necessarily equal to the actual number of
69/// columns.
70/// Each alias will be mapped to a column index depending on the step kind
71/// (Sponge or Round) that is currently being executed.
72#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
73pub enum ColumnAlias {
74 /// Hash identifier to distinguish inside the syscalls communication channel
75 HashIndex,
76 /// Block index inside the hash to enumerate preimage bytes
77 BlockIndex,
78 /// Hash step identifier to distinguish inside interstep communication
79 StepIndex,
80
81 Input(usize), // Curr[0..100) either ThetaStateA or SpongeOldState
82 Output(usize), // Next[0..100) either IotaStateG or SpongeXorState
83
84 ThetaShiftsC(usize), // Round Curr[100..180)
85 ThetaDenseC(usize), // Round Curr[180..200)
86 ThetaQuotientC(usize), // Round Curr[200..205)
87 ThetaRemainderC(usize), // Round Curr[205..225)
88 ThetaDenseRotC(usize), // Round Curr[225..245)
89 ThetaExpandRotC(usize), // Round Curr[245..265)
90 PiRhoShiftsE(usize), // Round Curr[265..665)
91 PiRhoDenseE(usize), // Round Curr[665..765)
92 PiRhoQuotientE(usize), // Round Curr[765..865)
93 PiRhoRemainderE(usize), // Round Curr[865..965)
94 PiRhoDenseRotE(usize), // Round Curr[965..1065)
95 PiRhoExpandRotE(usize), // Round Curr[1065..1165)
96 ChiShiftsB(usize), // Round Curr[1165..1565)
97 ChiShiftsSum(usize), // Round Curr[1565..1965)
98
99 SpongeNewState(usize), // Sponge Curr[100..200)
100 SpongeZeros(usize), // Sponge Curr[168..200)
101 SpongeBytes(usize), // Sponge Curr[200..400)
102 SpongeShifts(usize), // Sponge Curr[400..800)
103
104 RoundNumber, // Only nonzero when Selector(Flag::Round) = 1 : Round 0 | 1 ..=23
105 RoundConstants(usize), // Only nonzero when Selector(Flag::Round) = 1 : Round constants
106
107 PadLength, // Only nonzero when Selector(Flag::Pad) = 1 : Length 0 | 1 ..=136
108 TwoToPad, // Only nonzero when Selector(Flag::Pad) = 1 : 2^PadLength
109 PadSuffix(usize), // Only nonzero when Selector(Flag::Pad) = 1 : 5 field elements
110 PadBytesFlags(usize), // Only nonzero when Selector(Flag::Pad) = 1 : 136 boolean values
111}
112
113/// Variants of Keccak steps available for the interpreter.
114/// These selectors determine the specific behaviour so that Keccak steps
115/// can be split into different instances for folding
116#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, EnumIter, EnumCount)]
117pub enum Steps {
118 /// Current step performs a round of the permutation.
119 /// The round number stored in the Step is only used for the environment execution.
120 Round(u64),
121 /// Current step is a sponge
122 Sponge(Sponges),
123}
124/// Variants of Keccak sponges
125#[derive(
126 Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, EnumIter, EnumCount, Default,
127)]
128pub enum Sponges {
129 Absorb(Absorbs),
130 #[default]
131 Squeeze,
132}
133
134/// Order of absorb steps in the computation depending on the number of blocks to absorb
135#[derive(
136 Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, EnumIter, EnumCount, Default,
137)]
138pub enum Absorbs {
139 First, // Also known as the root absorb
140 Middle, // Any other absorb
141 Last, // Also known as the padding absorb
142 #[default]
143 Only, // Only one block to absorb (preimage data is less than 136 bytes), both root and pad
144}
145
146impl IntoIterator for Steps {
147 type Item = Steps;
148 type IntoIter = std::vec::IntoIter<Self::Item>;
149
150 /// Iterate over the instruction variants
151 fn into_iter(self) -> Self::IntoIter {
152 match self {
153 Steps::Round(_) => vec![Steps::Round(0)].into_iter(),
154 Steps::Sponge(_) => {
155 let mut iter_contents = Vec::with_capacity(Absorbs::COUNT + 1);
156 iter_contents
157 .extend(Absorbs::iter().map(|absorb| Steps::Sponge(Sponges::Absorb(absorb))));
158 iter_contents.push(Steps::Sponge(Sponges::Squeeze));
159 iter_contents.into_iter()
160 }
161 }
162 }
163}
164
165impl From<Steps> for usize {
166 /// Returns the index of the column corresponding to the given selector.
167 /// They are located at the end of the witness columns.
168 fn from(step: Steps) -> usize {
169 match step {
170 Round(_) => FLAG_ROUND_OFF,
171 Sponge(sponge) => match sponge {
172 Absorb(absorb) => match absorb {
173 First => FLAG_FST_OFF,
174 Middle => FLAG_MID_OFF,
175 Last => FLAG_LST_OFF,
176 Only => FLAG_ONE_OFF,
177 },
178 Squeeze => FLAG_SQUEEZE_OFF,
179 },
180 }
181 }
182}
183
184// The columns used by the Keccak circuit.
185// The Keccak circuit is split into two main modes: Round and Sponge (split into Root, Absorb, Pad, RootPad, Squeeze).
186// The columns are shared between the Sponge and Round steps
187// (the total number of columns refers to the maximum of columns used by each mode)
188// The hash, block, and step indices are shared between both modes.
189impl From<ColumnAlias> for usize {
190 /// Returns the witness column index for the given alias
191 fn from(alias: ColumnAlias) -> usize {
192 match alias {
193 ColumnAlias::HashIndex => STATUS_OFF,
194 ColumnAlias::BlockIndex => STATUS_OFF + 1,
195 ColumnAlias::StepIndex => STATUS_OFF + 2,
196
197 ColumnAlias::Input(i) => {
198 assert!(i < STATE_LEN);
199 CURR_OFF + i
200 }
201 ColumnAlias::Output(i) => {
202 assert!(i < STATE_LEN);
203 NEXT_OFF + i
204 }
205
206 ColumnAlias::ThetaShiftsC(i) => {
207 assert!(i < THETA_SHIFTS_C_LEN);
208 CURR_OFF + THETA_SHIFTS_C_OFF + i
209 }
210 ColumnAlias::ThetaDenseC(i) => {
211 assert!(i < THETA_DENSE_C_LEN);
212 CURR_OFF + THETA_DENSE_C_OFF + i
213 }
214 ColumnAlias::ThetaQuotientC(i) => {
215 assert!(i < THETA_QUOTIENT_C_LEN);
216 CURR_OFF + THETA_QUOTIENT_C_OFF + i
217 }
218 ColumnAlias::ThetaRemainderC(i) => {
219 assert!(i < THETA_REMAINDER_C_LEN);
220 CURR_OFF + THETA_REMAINDER_C_OFF + i
221 }
222 ColumnAlias::ThetaDenseRotC(i) => {
223 assert!(i < THETA_DENSE_ROT_C_LEN);
224 CURR_OFF + THETA_DENSE_ROT_C_OFF + i
225 }
226 ColumnAlias::ThetaExpandRotC(i) => {
227 assert!(i < THETA_EXPAND_ROT_C_LEN);
228 CURR_OFF + THETA_EXPAND_ROT_C_OFF + i
229 }
230 ColumnAlias::PiRhoShiftsE(i) => {
231 assert!(i < PIRHO_SHIFTS_E_LEN);
232 CURR_OFF + PIRHO_SHIFTS_E_OFF + i
233 }
234 ColumnAlias::PiRhoDenseE(i) => {
235 assert!(i < PIRHO_DENSE_E_LEN);
236 CURR_OFF + PIRHO_DENSE_E_OFF + i
237 }
238 ColumnAlias::PiRhoQuotientE(i) => {
239 assert!(i < PIRHO_QUOTIENT_E_LEN);
240 CURR_OFF + PIRHO_QUOTIENT_E_OFF + i
241 }
242 ColumnAlias::PiRhoRemainderE(i) => {
243 assert!(i < PIRHO_REMAINDER_E_LEN);
244 CURR_OFF + PIRHO_REMAINDER_E_OFF + i
245 }
246 ColumnAlias::PiRhoDenseRotE(i) => {
247 assert!(i < PIRHO_DENSE_ROT_E_LEN);
248 CURR_OFF + PIRHO_DENSE_ROT_E_OFF + i
249 }
250 ColumnAlias::PiRhoExpandRotE(i) => {
251 assert!(i < PIRHO_EXPAND_ROT_E_LEN);
252 CURR_OFF + PIRHO_EXPAND_ROT_E_OFF + i
253 }
254 ColumnAlias::ChiShiftsB(i) => {
255 assert!(i < CHI_SHIFTS_B_LEN);
256 CURR_OFF + CHI_SHIFTS_B_OFF + i
257 }
258 ColumnAlias::ChiShiftsSum(i) => {
259 assert!(i < CHI_SHIFTS_SUM_LEN);
260 CURR_OFF + CHI_SHIFTS_SUM_OFF + i
261 }
262
263 ColumnAlias::SpongeNewState(i) => {
264 assert!(i < SPONGE_NEW_STATE_LEN);
265 CURR_OFF + SPONGE_NEW_STATE_OFF + i
266 }
267 ColumnAlias::SpongeZeros(i) => {
268 assert!(i < SPONGE_ZEROS_LEN);
269 CURR_OFF + SPONGE_ZEROS_OFF + i
270 }
271 ColumnAlias::SpongeBytes(i) => {
272 assert!(i < SPONGE_BYTES_LEN);
273 CURR_OFF + SPONGE_BYTES_OFF + i
274 }
275 ColumnAlias::SpongeShifts(i) => {
276 assert!(i < SPONGE_SHIFTS_LEN);
277 CURR_OFF + SPONGE_SHIFTS_OFF + i
278 }
279
280 ColumnAlias::RoundNumber => RC_OFF,
281 ColumnAlias::RoundConstants(i) => {
282 assert!(i < ROUND_CONST_LEN);
283 RC_OFF + 1 + i
284 }
285
286 ColumnAlias::PadLength => PAD_FLAGS_OFF + PAD_LEN_OFF,
287 ColumnAlias::TwoToPad => PAD_FLAGS_OFF + PAD_TWO_OFF,
288 ColumnAlias::PadSuffix(i) => {
289 assert!(i < PAD_SUFFIX_LEN);
290 PAD_FLAGS_OFF + PAD_SUFFIX_OFF + i
291 }
292 ColumnAlias::PadBytesFlags(i) => {
293 assert!(i < PAD_BYTES_LEN);
294 PAD_FLAGS_OFF + PAD_BYTES_OFF + i
295 }
296 }
297 }
298}
299
300/// The witness columns used by the Keccak circuit.
301/// The Keccak circuit is split into two main modes: Sponge and Round.
302/// The columns are shared between the Sponge and Round steps.
303/// The hash and step indices are shared between both modes.
304/// The row is split into the following entries:
305/// - hash_index: Which hash this is inside the circuit
306/// - block_index: Which block this is inside the hash
307/// - step_index: Which step this is inside the hash
308/// - curr: Contains 1965 witnesses used in the current step including Input
309/// - next: Contains the 100 Output witnesses
310/// - round_flags: contain 5 elements with information about the current round step
311/// - pad_flags: PadLength, TwoToPad, PadBytesFlags, PadSuffix
312/// - mode_flags: what kind of mode is running: round, root, absorb, pad, rootpad, squeeze. Only 1 of them can be active.
313///
314/// Keccak Witness Columns: KeccakWitness.cols
315/// -------------------------------------------------------
316/// | 0 | 1 | 2 | 3..=1967 | 1968..=2067 | 2068..=2071 |
317/// -------------------------------------------------------
318/// 0 -> hash_index
319/// 1 -> block_index
320/// 2 -> step_index
321/// 3..=1967 -> curr
322/// 3 1967
323/// <--------------------------------if_round<---------------------------------->
324/// <-------------if_sponge-------------->
325/// 3 802
326/// -> SPONGE: | -> ROUND:
327/// -> 3..=102: Input == SpongeOldState | -> 3..=102: Input == ThetaStateA
328/// -> 103..=202: SpongeNewState | -> 103..=182: ThetaShiftsC
329/// : 170..=202 -> SpongeZeros | -> 183..=202: ThetaDenseC
330/// -> 203..=402: SpongeBytes | -> 203..=207: ThetaQuotientC
331/// -> 403..=802: SpongeShifts | -> 208..=227: ThetaRemainderC
332/// | -> 228..=247: ThetaDenseRotC
333/// | -> 248..=267: ThetaExpandRotC
334/// | -> 268..=667: PiRhoShiftsE
335/// | -> 668..=767: PiRhoDenseE
336/// | -> 768..=867: PiRhoQuotientE
337/// | -> 868..=967: PiRhoRemainderE
338/// | -> 968..=1067: PiRhoDenseRotE
339/// | -> 1068..=1167: PiRhoExpandRotE
340/// | -> 1168..=1567: ChiShiftsB
341/// | -> 1568..=1967: ChiShiftsSum
342/// 1968..=2067 -> next
343/// -> 1968..=2067: Output (if Round, then IotaStateG, if Sponge then SpongeXorState)
344///
345/// 2068..=2072 -> round_flags
346/// -> 2068: RoundNumber
347/// -> 2069..=2072: RoundConstants
348///
349/// 2073..=2078 -> selectors
350///
351/// 803..=945 -> pad_flags
352/// -> 803: PadLength
353/// -> 804: TwoToPad
354/// -> 805..=809: PadSuffix
355/// -> 810..=945: PadBytesFlags
356///
357pub type KeccakWitness<T> = Witness<N_ZKVM_KECCAK_REL_COLS, T>;
358
359/// IMPLEMENTATIONS FOR COLUMN ALIAS
360
361impl<T: Clone> Index<ColumnAlias> for KeccakWitness<T> {
362 type Output = T;
363
364 /// Map the column alias to the actual column index.
365 /// Note that the column index depends on the step kind (Sponge or Round).
366 /// For instance, the column 800 represents PadLength in the Sponge step, while it
367 /// is used by intermediary values when executing the Round step.
368 fn index(&self, index: ColumnAlias) -> &Self::Output {
369 &self.cols[usize::from(index)]
370 }
371}
372
373impl<T: Clone> IndexMut<ColumnAlias> for KeccakWitness<T> {
374 fn index_mut(&mut self, index: ColumnAlias) -> &mut Self::Output {
375 &mut self.cols[usize::from(index)]
376 }
377}
378
379impl ColumnIndexer<RelationColumnType> for ColumnAlias {
380 const N_COL: usize = N_ZKVM_KECCAK_REL_COLS + N_ZKVM_KECCAK_SEL_COLS;
381 fn to_column(self) -> Column<RelationColumnType> {
382 Column::Relation(RelationColumnType::Scratch(usize::from(self)))
383 }
384}
385
386// IMPLEMENTATIONS FOR SELECTOR
387
388impl<T: Clone> Index<Steps> for KeccakWitness<T> {
389 type Output = T;
390
391 /// Map the column alias to the actual column index.
392 /// Note that the column index depends on the step kind (Sponge or Round).
393 /// For instance, the column 800 represents PadLength in the Sponge step, while it
394 /// is used by intermediary values when executing the Round step.
395 /// The selector columns are located at the end of the witness relation columns.
396 fn index(&self, index: Steps) -> &Self::Output {
397 &self.cols[usize::from(index)]
398 }
399}
400
401impl<T: Clone> IndexMut<Steps> for KeccakWitness<T> {
402 fn index_mut(&mut self, index: Steps) -> &mut Self::Output {
403 &mut self.cols[usize::from(index)]
404 }
405}
406
407impl ColumnIndexer<usize> for Steps {
408 const N_COL: usize = N_ZKVM_KECCAK_REL_COLS + N_ZKVM_KECCAK_SEL_COLS;
409 fn to_column(self) -> Column<usize> {
410 Column::DynamicSelector(usize::from(self) - N_ZKVM_KECCAK_REL_COLS)
411 }
412}