kimchi/circuits/polynomials/keccak/
witness.rs

1//! Keccak witness computation
2
3use crate::{
4    auto_clone,
5    circuits::{
6        polynomials::keccak::{
7            constants::{
8                CAPACITY_IN_BYTES, DIM, KECCAK_COLS, QUARTERS, RATE_IN_BYTES, ROUNDS, STATE_LEN,
9            },
10            Keccak, OFF,
11        },
12        witness::{self, IndexCell, Variables, WitnessCell},
13    },
14    grid, variable_map,
15};
16use ark_ff::PrimeField;
17use core::array;
18use num_bigint::BigUint;
19
20pub(crate) const SPARSE_RC: [[u64; QUARTERS]; ROUNDS] = [
21    [
22        0x0000000000000001,
23        0x0000000000000000,
24        0x0000000000000000,
25        0x0000000000000000,
26    ],
27    [
28        0x1000000010000010,
29        0x0000000000000000,
30        0x0000000000000000,
31        0x0000000000000000,
32    ],
33    [
34        0x1000000010001010,
35        0x0000000000000000,
36        0x0000000000000000,
37        0x1000000000000000,
38    ],
39    [
40        0x1000000000000000,
41        0x1000000000000000,
42        0x0000000000000000,
43        0x1000000000000000,
44    ],
45    [
46        0x1000000010001011,
47        0x0000000000000000,
48        0x0000000000000000,
49        0x0000000000000000,
50    ],
51    [
52        0x0000000000000001,
53        0x1000000000000000,
54        0x0000000000000000,
55        0x0000000000000000,
56    ],
57    [
58        0x1000000010000001,
59        0x1000000000000000,
60        0x0000000000000000,
61        0x1000000000000000,
62    ],
63    [
64        0x1000000000001001,
65        0x0000000000000000,
66        0x0000000000000000,
67        0x1000000000000000,
68    ],
69    [
70        0x0000000010001010,
71        0x0000000000000000,
72        0x0000000000000000,
73        0x0000000000000000,
74    ],
75    [
76        0x0000000010001000,
77        0x0000000000000000,
78        0x0000000000000000,
79        0x0000000000000000,
80    ],
81    [
82        0x1000000000001001,
83        0x1000000000000000,
84        0x0000000000000000,
85        0x0000000000000000,
86    ],
87    [
88        0x0000000000001010,
89        0x1000000000000000,
90        0x0000000000000000,
91        0x0000000000000000,
92    ],
93    [
94        0x1000000010001011,
95        0x1000000000000000,
96        0x0000000000000000,
97        0x0000000000000000,
98    ],
99    [
100        0x0000000010001011,
101        0x0000000000000000,
102        0x0000000000000000,
103        0x1000000000000000,
104    ],
105    [
106        0x1000000010001001,
107        0x0000000000000000,
108        0x0000000000000000,
109        0x1000000000000000,
110    ],
111    [
112        0x1000000000000011,
113        0x0000000000000000,
114        0x0000000000000000,
115        0x1000000000000000,
116    ],
117    [
118        0x1000000000000010,
119        0x0000000000000000,
120        0x0000000000000000,
121        0x1000000000000000,
122    ],
123    [
124        0x0000000010000000,
125        0x0000000000000000,
126        0x0000000000000000,
127        0x1000000000000000,
128    ],
129    [
130        0x1000000000001010,
131        0x0000000000000000,
132        0x0000000000000000,
133        0x0000000000000000,
134    ],
135    [
136        0x0000000000001010,
137        0x1000000000000000,
138        0x0000000000000000,
139        0x1000000000000000,
140    ],
141    [
142        0x1000000010000001,
143        0x1000000000000000,
144        0x0000000000000000,
145        0x1000000000000000,
146    ],
147    [
148        0x1000000010000000,
149        0x0000000000000000,
150        0x0000000000000000,
151        0x1000000000000000,
152    ],
153    [
154        0x0000000000000001,
155        0x1000000000000000,
156        0x0000000000000000,
157        0x0000000000000000,
158    ],
159    [
160        0x1000000000001000,
161        0x1000000000000000,
162        0x0000000000000000,
163        0x1000000000000000,
164    ],
165];
166
167type Layout<F, const COLUMNS: usize> = Vec<Box<dyn WitnessCell<F, Vec<F>, COLUMNS>>>;
168
169fn layout_round<F: PrimeField>() -> [Layout<F, KECCAK_COLS>; 1] {
170    [vec![
171        IndexCell::create("state_a", 0, 100),
172        IndexCell::create("shifts_c", 100, 180),
173        IndexCell::create("dense_c", 180, 200),
174        IndexCell::create("quotient_c", 200, 205),
175        IndexCell::create("remainder_c", 205, 225),
176        IndexCell::create("dense_rot_c", 225, 245),
177        IndexCell::create("expand_rot_c", 245, 265),
178        IndexCell::create("shifts_e", 265, 665),
179        IndexCell::create("dense_e", 665, 765),
180        IndexCell::create("quotient_e", 765, 865),
181        IndexCell::create("remainder_e", 865, 965),
182        IndexCell::create("dense_rot_e", 965, 1065),
183        IndexCell::create("expand_rot_e", 1065, 1165),
184        IndexCell::create("shifts_b", 1165, 1565),
185        IndexCell::create("shifts_sum", 1565, 1965),
186    ]]
187}
188
189fn layout_sponge<F: PrimeField>() -> [Layout<F, KECCAK_COLS>; 1] {
190    [vec![
191        IndexCell::create("old_state", 0, 100),
192        IndexCell::create("new_state", 100, 200),
193        IndexCell::create("bytes", 200, 400),
194        IndexCell::create("shifts", 400, 800),
195    ]]
196}
197
198// Transforms a vector of u64 into a vector of field elements
199fn field<F: PrimeField>(input: &[u64]) -> Vec<F> {
200    input.iter().map(|x| F::from(*x)).collect::<Vec<F>>()
201}
202
203// Contains the quotient, remainder, bound, dense rotated as quarters of at most 16 bits each
204// Contains the expansion of the rotated word
205pub struct Rotation {
206    quotient: Vec<u64>,
207    remainder: Vec<u64>,
208    dense_rot: Vec<u64>,
209    expand_rot: Vec<u64>,
210}
211
212impl Rotation {
213    // On input the dense quarters of a word, rotate the word offset bits to the left
214    fn new(dense: &[u64], offset: u32) -> Self {
215        let word = Keccak::compose(dense);
216        let rem = word as u128 * 2u128.pow(offset) % 2u128.pow(64);
217        let quo = (word as u128) / 2u128.pow(64 - offset);
218        let rot = rem + quo;
219        assert!(rot as u64 == word.rotate_left(offset));
220
221        Self {
222            quotient: Keccak::decompose(quo as u64),
223            remainder: Keccak::decompose(rem as u64),
224            dense_rot: Keccak::decompose(rot as u64),
225            expand_rot: Keccak::decompose(rot as u64)
226                .iter()
227                .map(|x| Keccak::expand(*x))
228                .collect(),
229        }
230    }
231
232    // On input the dense quarters of many words, rotate the word offset bits to the left
233    fn many(words: &[u64], offsets: &[u32]) -> Self {
234        assert!(words.len() == QUARTERS * offsets.len());
235        let mut quotient = vec![];
236        let mut remainder = vec![];
237        let mut dense_rot = vec![];
238        let mut expand_rot = vec![];
239        for (word, offset) in words.chunks(QUARTERS).zip(offsets.iter()) {
240            let mut rot = Self::new(word, *offset);
241            quotient.append(&mut rot.quotient);
242            remainder.append(&mut rot.remainder);
243            dense_rot.append(&mut rot.dense_rot);
244            expand_rot.append(&mut rot.expand_rot);
245        }
246        Self {
247            quotient,
248            remainder,
249            dense_rot,
250            expand_rot,
251        }
252    }
253}
254
255/// Values involved in Theta permutation step
256pub struct Theta {
257    shifts_c: Vec<u64>,
258    dense_c: Vec<u64>,
259    quotient_c: Vec<u64>,
260    remainder_c: Vec<u64>,
261    dense_rot_c: Vec<u64>,
262    expand_rot_c: Vec<u64>,
263    state_e: Vec<u64>,
264}
265
266impl Theta {
267    pub fn create(state_a: &[u64]) -> Self {
268        let state_c = Self::compute_state_c(state_a);
269        let shifts_c = Keccak::shift(&state_c);
270        let dense_c = Keccak::collapse(&Keccak::reset(&shifts_c));
271        let rotation_c = Rotation::many(&dense_c, &[1; DIM]);
272        let state_d = Self::compute_state_d(&shifts_c, &rotation_c.expand_rot);
273        let state_e = Self::compute_state_e(state_a, &state_d);
274        let quotient_c = vec![
275            rotation_c.quotient[0],
276            rotation_c.quotient[4],
277            rotation_c.quotient[8],
278            rotation_c.quotient[12],
279            rotation_c.quotient[16],
280        ];
281        Self {
282            shifts_c,
283            dense_c,
284            quotient_c,
285            remainder_c: rotation_c.remainder,
286            dense_rot_c: rotation_c.dense_rot,
287            expand_rot_c: rotation_c.expand_rot,
288            state_e,
289        }
290    }
291
292    pub fn shifts_c(&self, i: usize, x: usize, q: usize) -> u64 {
293        let shifts_c = grid!(80, &self.shifts_c);
294        shifts_c(i, x, q)
295    }
296
297    pub fn dense_c(&self, x: usize, q: usize) -> u64 {
298        let dense_c = grid!(20, &self.dense_c);
299        dense_c(x, q)
300    }
301
302    pub fn quotient_c(&self, x: usize) -> u64 {
303        self.quotient_c[x]
304    }
305
306    pub fn remainder_c(&self, x: usize, q: usize) -> u64 {
307        let remainder_c = grid!(20, &self.remainder_c);
308        remainder_c(x, q)
309    }
310
311    pub fn dense_rot_c(&self, x: usize, q: usize) -> u64 {
312        let dense_rot_c = grid!(20, &self.dense_rot_c);
313        dense_rot_c(x, q)
314    }
315
316    pub fn expand_rot_c(&self, x: usize, q: usize) -> u64 {
317        let expand_rot_c = grid!(20, &self.expand_rot_c);
318        expand_rot_c(x, q)
319    }
320
321    pub fn state_e(&self) -> Vec<u64> {
322        self.state_e.clone()
323    }
324
325    fn compute_state_c(state_a: &[u64]) -> Vec<u64> {
326        let state_a = grid!(100, state_a);
327        let mut state_c = vec![];
328        for x in 0..DIM {
329            for q in 0..QUARTERS {
330                state_c.push(
331                    state_a(0, x, q)
332                        + state_a(1, x, q)
333                        + state_a(2, x, q)
334                        + state_a(3, x, q)
335                        + state_a(4, x, q),
336                );
337            }
338        }
339        state_c
340    }
341
342    fn compute_state_d(shifts_c: &[u64], expand_rot_c: &[u64]) -> Vec<u64> {
343        let shifts_c = grid!(20, shifts_c);
344        let expand_rot_c = grid!(20, expand_rot_c);
345        let mut state_d = vec![];
346        for x in 0..DIM {
347            for q in 0..QUARTERS {
348                state_d.push(shifts_c((x + DIM - 1) % DIM, q) + expand_rot_c((x + 1) % DIM, q));
349            }
350        }
351        state_d
352    }
353
354    fn compute_state_e(state_a: &[u64], state_d: &[u64]) -> Vec<u64> {
355        let state_a = grid!(100, state_a);
356        let state_d = grid!(20, state_d);
357        let mut state_e = vec![];
358        for y in 0..DIM {
359            for x in 0..DIM {
360                for q in 0..QUARTERS {
361                    state_e.push(state_a(y, x, q) + state_d(x, q));
362                }
363            }
364        }
365        state_e
366    }
367}
368
369/// Values involved in PiRho permutation step
370pub struct PiRho {
371    shifts_e: Vec<u64>,
372    dense_e: Vec<u64>,
373    quotient_e: Vec<u64>,
374    remainder_e: Vec<u64>,
375    dense_rot_e: Vec<u64>,
376    expand_rot_e: Vec<u64>,
377    state_b: Vec<u64>,
378}
379
380impl PiRho {
381    pub fn create(state_e: &[u64]) -> Self {
382        let shifts_e = Keccak::shift(state_e);
383        let dense_e = Keccak::collapse(&Keccak::reset(&shifts_e));
384        let rotation_e = Rotation::many(
385            &dense_e,
386            &OFF.iter()
387                .flatten()
388                .map(|x| *x as u32)
389                .collect::<Vec<u32>>(),
390        );
391
392        let mut state_b = vec![vec![vec![0; QUARTERS]; DIM]; DIM];
393        let aux = grid!(100, rotation_e.expand_rot);
394        for y in 0..DIM {
395            for x in 0..DIM {
396                #[allow(clippy::needless_range_loop)]
397                for q in 0..QUARTERS {
398                    state_b[(2 * x + 3 * y) % DIM][y][q] = aux(y, x, q);
399                }
400            }
401        }
402        let state_b = state_b.iter().flatten().flatten().copied().collect();
403
404        Self {
405            shifts_e,
406            dense_e,
407            quotient_e: rotation_e.quotient,
408            remainder_e: rotation_e.remainder,
409            dense_rot_e: rotation_e.dense_rot,
410            expand_rot_e: rotation_e.expand_rot,
411            state_b,
412        }
413    }
414
415    pub fn shifts_e(&self, i: usize, y: usize, x: usize, q: usize) -> u64 {
416        let shifts_e = grid!(400, &self.shifts_e);
417        shifts_e(i, y, x, q)
418    }
419
420    pub fn dense_e(&self, y: usize, x: usize, q: usize) -> u64 {
421        let dense_e = grid!(100, &self.dense_e);
422        dense_e(y, x, q)
423    }
424
425    pub fn quotient_e(&self, y: usize, x: usize, q: usize) -> u64 {
426        let quotient_e = grid!(100, &self.quotient_e);
427        quotient_e(y, x, q)
428    }
429
430    pub fn remainder_e(&self, y: usize, x: usize, q: usize) -> u64 {
431        let remainder_e = grid!(100, &self.remainder_e);
432        remainder_e(y, x, q)
433    }
434
435    pub fn dense_rot_e(&self, y: usize, x: usize, q: usize) -> u64 {
436        let dense_rot_e = grid!(100, &self.dense_rot_e);
437        dense_rot_e(y, x, q)
438    }
439
440    pub fn expand_rot_e(&self, y: usize, x: usize, q: usize) -> u64 {
441        let expand_rot_e = grid!(100, &self.expand_rot_e);
442        expand_rot_e(y, x, q)
443    }
444
445    pub fn state_b(&self) -> Vec<u64> {
446        self.state_b.clone()
447    }
448}
449
450/// Values involved in Chi permutation step
451pub struct Chi {
452    shifts_b: Vec<u64>,
453    shifts_sum: Vec<u64>,
454    state_f: Vec<u64>,
455}
456
457impl Chi {
458    pub fn create(state_b: &[u64]) -> Self {
459        let shifts_b = Keccak::shift(state_b);
460        let shiftsb = grid!(400, shifts_b);
461        let mut sum = vec![];
462        for y in 0..DIM {
463            for x in 0..DIM {
464                for q in 0..QUARTERS {
465                    let not = 0x1111111111111111u64 - shiftsb(0, y, (x + 1) % DIM, q);
466                    sum.push(not + shiftsb(0, y, (x + 2) % DIM, q));
467                }
468            }
469        }
470        let shifts_sum = Keccak::shift(&sum);
471        let shiftsum = grid!(400, shifts_sum);
472        let mut state_f = vec![];
473        for y in 0..DIM {
474            for x in 0..DIM {
475                for q in 0..QUARTERS {
476                    let and = shiftsum(1, y, x, q);
477                    state_f.push(shiftsb(0, y, x, q) + and);
478                }
479            }
480        }
481
482        Self {
483            shifts_b,
484            shifts_sum,
485            state_f,
486        }
487    }
488
489    pub fn shifts_b(&self, i: usize, y: usize, x: usize, q: usize) -> u64 {
490        let shifts_b = grid!(400, &self.shifts_b);
491        shifts_b(i, y, x, q)
492    }
493
494    pub fn shifts_sum(&self, i: usize, y: usize, x: usize, q: usize) -> u64 {
495        let shifts_sum = grid!(400, &self.shifts_sum);
496        shifts_sum(i, y, x, q)
497    }
498
499    pub fn state_f(&self) -> Vec<u64> {
500        self.state_f.clone()
501    }
502}
503
504/// Values involved in Iota permutation step
505pub struct Iota {
506    state_g: Vec<u64>,
507    round_constants: [u64; QUARTERS],
508}
509
510impl Iota {
511    pub fn create(state_f: &[u64], round: usize) -> Self {
512        let round_constants = SPARSE_RC[round];
513        let mut state_g = state_f.to_vec();
514        for (i, c) in round_constants.iter().enumerate() {
515            state_g[i] = state_f[i] + *c;
516        }
517        Self {
518            state_g,
519            round_constants,
520        }
521    }
522
523    pub fn state_g(&self) -> Vec<u64> {
524        self.state_g.clone()
525    }
526
527    pub fn round_constants(&self, i: usize) -> u64 {
528        self.round_constants[i]
529    }
530}
531
532/// Creates a witness for the Keccak hash function
533/// Input:
534/// - message: the message to be hashed
535///
536/// Note:
537///   Requires at least one more row after the keccak gadget so that
538///   constraints can access the next row in the squeeze
539pub fn extend_keccak_witness<F: PrimeField>(witness: &mut [Vec<F>; KECCAK_COLS], message: BigUint) {
540    let padded = Keccak::pad(&message.to_bytes_be());
541    let chunks = padded.chunks(RATE_IN_BYTES);
542
543    // The number of rows that need to be added to the witness correspond to
544    // - Absorb phase:
545    //      - 1 per block for the sponge row
546    //      - 24 for the rounds
547    // - Squeeze phase:
548    //      - 1 for the final sponge row
549    let rows: usize = chunks.len() * (ROUNDS + 1) + 1;
550
551    let mut keccak_witness = array::from_fn(|_| vec![F::zero(); rows]);
552
553    // Absorb phase
554    let mut row = 0;
555    let mut state = vec![0; QUARTERS * DIM * DIM];
556    for chunk in chunks {
557        let mut block = chunk.to_vec();
558        // Pad the block until reaching 200 bytes
559        block.append(&mut vec![0; CAPACITY_IN_BYTES]);
560        let new_state = Keccak::expand_state(&block);
561        auto_clone!(new_state);
562        let shifts = Keccak::shift(&new_state());
563        let bytes = block.iter().map(|b| *b as u64).collect::<Vec<u64>>();
564
565        // Initialize the absorb sponge row
566        witness::init(
567            &mut keccak_witness,
568            row,
569            &layout_sponge(),
570            &variable_map!["old_state" => field(&state), "new_state" => field(&new_state()), "bytes" => field(&bytes), "shifts" => field(&shifts)],
571        );
572        row += 1;
573
574        let xor_state = state
575            .iter()
576            .zip(new_state())
577            .map(|(x, y)| x + y)
578            .collect::<Vec<u64>>();
579
580        let mut ini_state = xor_state.clone();
581
582        for round in 0..ROUNDS {
583            // Theta
584            let theta = Theta::create(&ini_state);
585
586            // PiRho
587            let pirho = PiRho::create(&theta.state_e);
588
589            // Chi
590            let chi = Chi::create(&pirho.state_b);
591
592            // Iota
593            let iota = Iota::create(&chi.state_f, round);
594
595            // Initialize the round row
596            witness::init(
597                &mut keccak_witness,
598                row,
599                &layout_round(),
600                &variable_map![
601                "state_a" => field(&ini_state),
602                "shifts_c" => field(&theta.shifts_c),
603                "dense_c" => field(&theta.dense_c),
604                "quotient_c" => field(&theta.quotient_c),
605                "remainder_c" => field(&theta.remainder_c),
606                "dense_rot_c" => field(&theta.dense_rot_c),
607                "expand_rot_c" => field(&theta.expand_rot_c),
608                "shifts_e" => field(&pirho.shifts_e),
609                "dense_e" => field(&pirho.dense_e),
610                "quotient_e" => field(&pirho.quotient_e),
611                "remainder_e" => field(&pirho.remainder_e),
612                "dense_rot_e" => field(&pirho.dense_rot_e),
613                "expand_rot_e" => field(&pirho.expand_rot_e),
614                "shifts_b" => field(&chi.shifts_b),
615                "shifts_sum" => field(&chi.shifts_sum)
616                ],
617            );
618            row += 1;
619            ini_state = iota.state_g;
620        }
621        // update state after rounds
622        state = ini_state;
623    }
624
625    // Squeeze phase
626
627    let new_state = vec![0; STATE_LEN];
628    let shifts = Keccak::shift(&state);
629    let dense = Keccak::collapse(&Keccak::reset(&shifts));
630    let bytes = Keccak::bytestring(&dense);
631
632    // Initialize the squeeze sponge row
633    witness::init(
634        &mut keccak_witness,
635        row,
636        &layout_sponge(),
637        &variable_map!["old_state" => field(&state), "new_state" => field(&new_state), "bytes" => field(&bytes), "shifts" => field(&shifts)],
638    );
639
640    for col in 0..KECCAK_COLS {
641        witness[col].extend(keccak_witness[col].iter());
642    }
643}
644
645#[cfg(test)]
646mod tests {
647    use super::*;
648    use crate::circuits::polynomials::keccak::RC;
649
650    #[test]
651    fn test_sparse_round_constants() {
652        for round in 0..ROUNDS {
653            let round_constants = Keccak::sparse(RC[round]);
654            for (i, rc) in round_constants.iter().enumerate().take(QUARTERS) {
655                assert_eq!(*rc, SPARSE_RC[round][i]);
656            }
657        }
658    }
659}
kimchi/circuits/polynomials/keccak/witness.rs

kimchi/circuits/polynomials/keccak/
witness.rs
