p3_poseidon2_air/
columns.rs

Help
1use core::borrow::{Borrow, BorrowMut};
2use core::mem::size_of;
3
4/// Columns for a Poseidon2 AIR which computes one permutation per row.
5///
6/// The columns of the STARK are divided into the three different round sections of the Poseidon2
7/// Permutation: beginning full rounds, partial rounds, and ending full rounds. For the full
8/// rounds we store an [`SBox`] columnset for each state variable, and for the partial rounds we
9/// store only for the first state variable. Because the matrix multiplications are linear
10/// functions, we need only keep auxiliary columns for the S-box computations.
11#[repr(C)]
12pub struct Poseidon2Cols<
13    T,
14    const WIDTH: usize,
15    const SBOX_DEGREE: u64,
16    const SBOX_REGISTERS: usize,
17    const HALF_FULL_ROUNDS: usize,
18    const PARTIAL_ROUNDS: usize,
19> {
20    pub export: T,
21
22    pub inputs: [T; WIDTH],
23
24    /// Beginning Full Rounds
25    pub beginning_full_rounds: [FullRound<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>; HALF_FULL_ROUNDS],
26
27    /// Partial Rounds
28    pub partial_rounds: [PartialRound<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>; PARTIAL_ROUNDS],
29
30    /// Ending Full Rounds
31    pub ending_full_rounds: [FullRound<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>; HALF_FULL_ROUNDS],
32}
33
34/// Full round columns.
35#[repr(C)]
36pub struct FullRound<T, const WIDTH: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize> {
37    /// Possible intermediate results within each S-box.
38    pub sbox: [SBox<T, SBOX_DEGREE, SBOX_REGISTERS>; WIDTH],
39    /// The post-state, i.e. the entire layer after this full round.
40    pub post: [T; WIDTH],
41}
42
43/// Partial round columns.
44#[repr(C)]
45pub struct PartialRound<T, const WIDTH: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize>
46{
47    /// Possible intermediate results within the S-box.
48    pub sbox: SBox<T, SBOX_DEGREE, SBOX_REGISTERS>,
49    /// The output of the S-box.
50    pub post_sbox: T,
51}
52
53/// Possible intermediate results within an S-box.
54///
55/// Use this column-set for an S-box that can be computed with `REGISTERS`-many intermediate results
56/// (not counting the final output). The S-box is checked to ensure that `REGISTERS` is the optimal
57/// number of registers for the given `DEGREE` for the degrees given in the Poseidon2 paper:
58/// `3`, `5`, `7`, and `11`. See `eval_sbox` for more information.
59#[repr(C)]
60pub struct SBox<T, const DEGREE: u64, const REGISTERS: usize>(pub [T; REGISTERS]);
61
62pub const fn num_cols<
63    const WIDTH: usize,
64    const SBOX_DEGREE: u64,
65    const SBOX_REGISTERS: usize,
66    const HALF_FULL_ROUNDS: usize,
67    const PARTIAL_ROUNDS: usize,
68>() -> usize {
69    size_of::<Poseidon2Cols<u8, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>>(
70    )
71}
72
73pub const fn make_col_map<
74    const WIDTH: usize,
75    const SBOX_DEGREE: u64,
76    const SBOX_REGISTERS: usize,
77    const HALF_FULL_ROUNDS: usize,
78    const PARTIAL_ROUNDS: usize,
79>() -> Poseidon2Cols<usize, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS> {
80    todo!()
81    // let indices_arr = indices_arr::<
82    //     { num_cols::<WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>() },
83    // >();
84    // unsafe {
85    //     transmute::<
86    //         [usize;
87    //             num_cols::<WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>()],
88    //         Poseidon2Cols<
89    //             usize,
90    //             WIDTH,
91    //             SBOX_DEGREE,
92    //             SBOX_REGISTERS,
93    //             HALF_FULL_ROUNDS,
94    //             PARTIAL_ROUNDS,
95    //         >,
96    //     >(indices_arr)
97    // }
98}
99
100impl<
101        T,
102        const WIDTH: usize,
103        const SBOX_DEGREE: u64,
104        const SBOX_REGISTERS: usize,
105        const HALF_FULL_ROUNDS: usize,
106        const PARTIAL_ROUNDS: usize,
107    > Borrow<Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>>
108    for [T]
109{
110    fn borrow(
111        &self,
112    ) -> &Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>
113    {
114        // debug_assert_eq!(self.len(), NUM_COLS);
115        let (prefix, shorts, suffix) = unsafe {
116            self.align_to::<Poseidon2Cols<
117                T,
118                WIDTH,
119                SBOX_DEGREE,
120                SBOX_REGISTERS,
121                HALF_FULL_ROUNDS,
122                PARTIAL_ROUNDS,
123            >>()
124        };
125        debug_assert!(prefix.is_empty(), "Alignment should match");
126        debug_assert!(suffix.is_empty(), "Alignment should match");
127        debug_assert_eq!(shorts.len(), 1);
128        &shorts[0]
129    }
130}
131
132impl<
133        T,
134        const WIDTH: usize,
135        const SBOX_DEGREE: u64,
136        const SBOX_REGISTERS: usize,
137        const HALF_FULL_ROUNDS: usize,
138        const PARTIAL_ROUNDS: usize,
139    >
140    BorrowMut<
141        Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
142    > for [T]
143{
144    fn borrow_mut(
145        &mut self,
146    ) -> &mut Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>
147    {
148        // debug_assert_eq!(self.len(), NUM_COLS);
149        let (prefix, shorts, suffix) = unsafe {
150            self.align_to_mut::<Poseidon2Cols<
151                T,
152                WIDTH,
153                SBOX_DEGREE,
154                SBOX_REGISTERS,
155                HALF_FULL_ROUNDS,
156                PARTIAL_ROUNDS,
157            >>()
158        };
159        debug_assert!(prefix.is_empty(), "Alignment should match");
160        debug_assert!(suffix.is_empty(), "Alignment should match");
161        debug_assert_eq!(shorts.len(), 1);
162        &mut shorts[0]
163    }
164}
p3_poseidon2_air/columns.rs

p3_poseidon2_air/
columns.rs
