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                for q in 0..QUARTERS {
397                    state_b[(2 * x + 3 * y) % DIM][y][q] = aux(y, x, q);
398                }
399            }
400        }
401        let state_b = state_b.iter().flatten().flatten().copied().collect();
402
403        Self {
404            shifts_e,
405            dense_e,
406            quotient_e: rotation_e.quotient,
407            remainder_e: rotation_e.remainder,
408            dense_rot_e: rotation_e.dense_rot,
409            expand_rot_e: rotation_e.expand_rot,
410            state_b,
411        }
412    }
413
414    pub fn shifts_e(&self, i: usize, y: usize, x: usize, q: usize) -> u64 {
415        let shifts_e = grid!(400, &self.shifts_e);
416        shifts_e(i, y, x, q)
417    }
418
419    pub fn dense_e(&self, y: usize, x: usize, q: usize) -> u64 {
420        let dense_e = grid!(100, &self.dense_e);
421        dense_e(y, x, q)
422    }
423
424    pub fn quotient_e(&self, y: usize, x: usize, q: usize) -> u64 {
425        let quotient_e = grid!(100, &self.quotient_e);
426        quotient_e(y, x, q)
427    }
428
429    pub fn remainder_e(&self, y: usize, x: usize, q: usize) -> u64 {
430        let remainder_e = grid!(100, &self.remainder_e);
431        remainder_e(y, x, q)
432    }
433
434    pub fn dense_rot_e(&self, y: usize, x: usize, q: usize) -> u64 {
435        let dense_rot_e = grid!(100, &self.dense_rot_e);
436        dense_rot_e(y, x, q)
437    }
438
439    pub fn expand_rot_e(&self, y: usize, x: usize, q: usize) -> u64 {
440        let expand_rot_e = grid!(100, &self.expand_rot_e);
441        expand_rot_e(y, x, q)
442    }
443
444    pub fn state_b(&self) -> Vec<u64> {
445        self.state_b.clone()
446    }
447}
448
449/// Values involved in Chi permutation step
450pub struct Chi {
451    shifts_b: Vec<u64>,
452    shifts_sum: Vec<u64>,
453    state_f: Vec<u64>,
454}
455
456impl Chi {
457    pub fn create(state_b: &[u64]) -> Self {
458        let shifts_b = Keccak::shift(state_b);
459        let shiftsb = grid!(400, shifts_b);
460        let mut sum = vec![];
461        for y in 0..DIM {
462            for x in 0..DIM {
463                for q in 0..QUARTERS {
464                    let not = 0x1111111111111111u64 - shiftsb(0, y, (x + 1) % DIM, q);
465                    sum.push(not + shiftsb(0, y, (x + 2) % DIM, q));
466                }
467            }
468        }
469        let shifts_sum = Keccak::shift(&sum);
470        let shiftsum = grid!(400, shifts_sum);
471        let mut state_f = vec![];
472        for y in 0..DIM {
473            for x in 0..DIM {
474                for q in 0..QUARTERS {
475                    let and = shiftsum(1, y, x, q);
476                    state_f.push(shiftsb(0, y, x, q) + and);
477                }
478            }
479        }
480
481        Self {
482            shifts_b,
483            shifts_sum,
484            state_f,
485        }
486    }
487
488    pub fn shifts_b(&self, i: usize, y: usize, x: usize, q: usize) -> u64 {
489        let shifts_b = grid!(400, &self.shifts_b);
490        shifts_b(i, y, x, q)
491    }
492
493    pub fn shifts_sum(&self, i: usize, y: usize, x: usize, q: usize) -> u64 {
494        let shifts_sum = grid!(400, &self.shifts_sum);
495        shifts_sum(i, y, x, q)
496    }
497
498    pub fn state_f(&self) -> Vec<u64> {
499        self.state_f.clone()
500    }
501}
502
503/// Values involved in Iota permutation step
504pub struct Iota {
505    state_g: Vec<u64>,
506    round_constants: [u64; QUARTERS],
507}
508
509impl Iota {
510    pub fn create(state_f: &[u64], round: usize) -> Self {
511        let round_constants = SPARSE_RC[round];
512        let mut state_g = state_f.to_vec();
513        for (i, c) in round_constants.iter().enumerate() {
514            state_g[i] = state_f[i] + *c;
515        }
516        Self {
517            state_g,
518            round_constants,
519        }
520    }
521
522    pub fn state_g(&self) -> Vec<u64> {
523        self.state_g.clone()
524    }
525
526    pub fn round_constants(&self, i: usize) -> u64 {
527        self.round_constants[i]
528    }
529}
530
531/// Creates a witness for the Keccak hash function
532/// Input:
533/// - message: the message to be hashed
534///
535/// Note:
536///   Requires at least one more row after the keccak gadget so that
537///   constraints can access the next row in the squeeze
538pub fn extend_keccak_witness<F: PrimeField>(witness: &mut [Vec<F>; KECCAK_COLS], message: BigUint) {
539    let padded = Keccak::pad(&message.to_bytes_be());
540    let chunks = padded.chunks(RATE_IN_BYTES);
541
542    // The number of rows that need to be added to the witness correspond to
543    // - Absorb phase:
544    //      - 1 per block for the sponge row
545    //      - 24 for the rounds
546    // - Squeeze phase:
547    //      - 1 for the final sponge row
548    let rows: usize = chunks.len() * (ROUNDS + 1) + 1;
549
550    let mut keccak_witness = array::from_fn(|_| vec![F::zero(); rows]);
551
552    // Absorb phase
553    let mut row = 0;
554    let mut state = vec![0; QUARTERS * DIM * DIM];
555    for chunk in chunks {
556        let mut block = chunk.to_vec();
557        // Pad the block until reaching 200 bytes
558        block.append(&mut vec![0; CAPACITY_IN_BYTES]);
559        let new_state = Keccak::expand_state(&block);
560        auto_clone!(new_state);
561        let shifts = Keccak::shift(&new_state());
562        let bytes = block.iter().map(|b| *b as u64).collect::<Vec<u64>>();
563
564        // Initialize the absorb sponge row
565        witness::init(
566            &mut keccak_witness,
567            row,
568            &layout_sponge(),
569            &variable_map!["old_state" => field(&state), "new_state" => field(&new_state()), "bytes" => field(&bytes), "shifts" => field(&shifts)],
570        );
571        row += 1;
572
573        let xor_state = state
574            .iter()
575            .zip(new_state())
576            .map(|(x, y)| x + y)
577            .collect::<Vec<u64>>();
578
579        let mut ini_state = xor_state.clone();
580
581        for round in 0..ROUNDS {
582            // Theta
583            let theta = Theta::create(&ini_state);
584
585            // PiRho
586            let pirho = PiRho::create(&theta.state_e);
587
588            // Chi
589            let chi = Chi::create(&pirho.state_b);
590
591            // Iota
592            let iota = Iota::create(&chi.state_f, round);
593
594            // Initialize the round row
595            witness::init(
596                &mut keccak_witness,
597                row,
598                &layout_round(),
599                &variable_map![
600                "state_a" => field(&ini_state),
601                "shifts_c" => field(&theta.shifts_c),
602                "dense_c" => field(&theta.dense_c),
603                "quotient_c" => field(&theta.quotient_c),
604                "remainder_c" => field(&theta.remainder_c),
605                "dense_rot_c" => field(&theta.dense_rot_c),
606                "expand_rot_c" => field(&theta.expand_rot_c),
607                "shifts_e" => field(&pirho.shifts_e),
608                "dense_e" => field(&pirho.dense_e),
609                "quotient_e" => field(&pirho.quotient_e),
610                "remainder_e" => field(&pirho.remainder_e),
611                "dense_rot_e" => field(&pirho.dense_rot_e),
612                "expand_rot_e" => field(&pirho.expand_rot_e),
613                "shifts_b" => field(&chi.shifts_b),
614                "shifts_sum" => field(&chi.shifts_sum)
615                ],
616            );
617            row += 1;
618            ini_state = iota.state_g;
619        }
620        // update state after rounds
621        state = ini_state;
622    }
623
624    // Squeeze phase
625
626    let new_state = vec![0; STATE_LEN];
627    let shifts = Keccak::shift(&state);
628    let dense = Keccak::collapse(&Keccak::reset(&shifts));
629    let bytes = Keccak::bytestring(&dense);
630
631    // Initialize the squeeze sponge row
632    witness::init(
633        &mut keccak_witness,
634        row,
635        &layout_sponge(),
636        &variable_map!["old_state" => field(&state), "new_state" => field(&new_state), "bytes" => field(&bytes), "shifts" => field(&shifts)],
637    );
638
639    for col in 0..KECCAK_COLS {
640        witness[col].extend(keccak_witness[col].iter());
641    }
642}
643
644#[cfg(test)]
645mod tests {
646    use super::*;
647    use crate::circuits::polynomials::keccak::RC;
648
649    #[test]
650    fn test_sparse_round_constants() {
651        for round in 0..ROUNDS {
652            let round_constants = Keccak::sparse(RC[round]);
653            for (i, rc) in round_constants.iter().enumerate().take(QUARTERS) {
654                assert_eq!(*rc, SPARSE_RC[round][i]);
655            }
656        }
657    }
658}
kimchi/circuits/polynomials/keccak/witness.rs

kimchi/circuits/polynomials/keccak/
witness.rs
