p3_keccak_air/
columns.rs

Help
1use core::borrow::{Borrow, BorrowMut};
2use core::mem::{size_of, transmute};
3
4use p3_util::indices_arr;
5
6use crate::constants::R;
7use crate::{NUM_ROUNDS, RATE_LIMBS, U64_LIMBS};
8
9/// Note: The ordering of each array is based on the input mapping. As the spec says,
10///
11/// > The mapping between the bits of s and those of a is `s[w(5y + x) + z] = a[x][y][z]`.
12///
13/// Thus, for example, `a_prime` is stored in `y, x, z` order. This departs from the more common
14/// convention of `x, y, z` order, but it has the benefit that input lists map to AIR columns in a
15/// nicer way.
16#[derive(Debug)]
17#[repr(C)]
18pub struct KeccakCols<T> {
19    /// The `i`th value is set to 1 if we are in the `i`th round, otherwise 0.
20    pub step_flags: [T; NUM_ROUNDS],
21
22    /// A register which indicates if a row should be exported, i.e. included in a multiset equality
23    /// argument. Should be 1 only for certain rows which are final steps, i.e. with
24    /// `step_flags[23] = 1`.
25    pub export: T,
26
27    /// Permutation inputs, stored in y-major order.
28    pub preimage: [[[T; U64_LIMBS]; 5]; 5],
29
30    pub a: [[[T; U64_LIMBS]; 5]; 5],
31
32    /// ```ignore
33    /// C[x] = xor(A[x, 0], A[x, 1], A[x, 2], A[x, 3], A[x, 4])
34    /// ```
35    pub c: [[T; 64]; 5],
36
37    /// ```ignore
38    /// C'[x, z] = xor(C[x, z], C[x - 1, z], C[x + 1, z - 1])
39    /// ```
40    pub c_prime: [[T; 64]; 5],
41
42    // Note: D is inlined, not stored in the witness.
43    /// ```ignore
44    /// A'[x, y] = xor(A[x, y], D[x])
45    ///          = xor(A[x, y], C[x - 1], ROT(C[x + 1], 1))
46    /// ```
47    pub a_prime: [[[T; 64]; 5]; 5],
48
49    /// ```ignore
50    /// A''[x, y] = xor(B[x, y], andn(B[x + 1, y], B[x + 2, y])).
51    /// ```
52    pub a_prime_prime: [[[T; U64_LIMBS]; 5]; 5],
53
54    /// The bits of `A''[0, 0]`.
55    pub a_prime_prime_0_0_bits: [T; 64],
56
57    /// ```ignore
58    /// A'''[0, 0, z] = A''[0, 0, z] ^ RC[k, z]
59    /// ```
60    pub a_prime_prime_prime_0_0_limbs: [T; U64_LIMBS],
61}
62
63impl<T: Copy> KeccakCols<T> {
64    pub fn b(&self, x: usize, y: usize, z: usize) -> T {
65        debug_assert!(x < 5);
66        debug_assert!(y < 5);
67        debug_assert!(z < 64);
68
69        // B is just a rotation of A', so these are aliases for A' registers.
70        // From the spec,
71        //     B[y, (2x + 3y) % 5] = ROT(A'[x, y], r[x, y])
72        // So,
73        //     B[x, y] = f((x + 3y) % 5, x)
74        // where f(a, b) = ROT(A'[a, b], r[a, b])
75        let a = (x + 3 * y) % 5;
76        let b = x;
77        let rot = R[a][b] as usize;
78        self.a_prime[b][a][(z + 64 - rot) % 64]
79    }
80
81    pub fn a_prime_prime_prime(&self, y: usize, x: usize, limb: usize) -> T {
82        debug_assert!(y < 5);
83        debug_assert!(x < 5);
84        debug_assert!(limb < U64_LIMBS);
85
86        if y == 0 && x == 0 {
87            self.a_prime_prime_prime_0_0_limbs[limb]
88        } else {
89            self.a_prime_prime[y][x][limb]
90        }
91    }
92}
93
94pub fn input_limb(i: usize) -> usize {
95    debug_assert!(i < RATE_LIMBS);
96
97    let i_u64 = i / U64_LIMBS;
98    let limb_index = i % U64_LIMBS;
99
100    // The 5x5 state is treated as y-major, as per the Keccak spec.
101    let y = i_u64 / 5;
102    let x = i_u64 % 5;
103
104    KECCAK_COL_MAP.preimage[y][x][limb_index]
105}
106
107pub fn output_limb(i: usize) -> usize {
108    debug_assert!(i < RATE_LIMBS);
109
110    let i_u64 = i / U64_LIMBS;
111    let limb_index = i % U64_LIMBS;
112
113    // The 5x5 state is treated as y-major, as per the Keccak spec.
114    let y = i_u64 / 5;
115    let x = i_u64 % 5;
116
117    KECCAK_COL_MAP.a_prime_prime_prime(y, x, limb_index)
118}
119
120pub const NUM_KECCAK_COLS: usize = size_of::<KeccakCols<u8>>();
121pub(crate) const KECCAK_COL_MAP: KeccakCols<usize> = make_col_map();
122
123const fn make_col_map() -> KeccakCols<usize> {
124    let indices_arr = indices_arr::<NUM_KECCAK_COLS>();
125    unsafe { transmute::<[usize; NUM_KECCAK_COLS], KeccakCols<usize>>(indices_arr) }
126}
127
128impl<T> Borrow<KeccakCols<T>> for [T] {
129    fn borrow(&self) -> &KeccakCols<T> {
130        debug_assert_eq!(self.len(), NUM_KECCAK_COLS);
131        let (prefix, shorts, suffix) = unsafe { self.align_to::<KeccakCols<T>>() };
132        debug_assert!(prefix.is_empty(), "Alignment should match");
133        debug_assert!(suffix.is_empty(), "Alignment should match");
134        debug_assert_eq!(shorts.len(), 1);
135        &shorts[0]
136    }
137}
138
139impl<T> BorrowMut<KeccakCols<T>> for [T] {
140    fn borrow_mut(&mut self) -> &mut KeccakCols<T> {
141        debug_assert_eq!(self.len(), NUM_KECCAK_COLS);
142        let (prefix, shorts, suffix) = unsafe { self.align_to_mut::<KeccakCols<T>>() };
143        debug_assert!(prefix.is_empty(), "Alignment should match");
144        debug_assert!(suffix.is_empty(), "Alignment should match");
145        debug_assert_eq!(shorts.len(), 1);
146        &mut shorts[0]
147    }
148}
p3_keccak_air/columns.rs

p3_keccak_air/
columns.rs
