p3_poseidon2_air/
columns.rs

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
73impl<
74    T,
75    const WIDTH: usize,
76    const SBOX_DEGREE: u64,
77    const SBOX_REGISTERS: usize,
78    const HALF_FULL_ROUNDS: usize,
79    const PARTIAL_ROUNDS: usize,
80> Borrow<Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>>
81    for [T]
82{
83    fn borrow(
84        &self,
85    ) -> &Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>
86    {
87        // debug_assert_eq!(self.len(), NUM_COLS);
88        let (prefix, shorts, suffix) = unsafe {
89            self.align_to::<Poseidon2Cols<
90                T,
91                WIDTH,
92                SBOX_DEGREE,
93                SBOX_REGISTERS,
94                HALF_FULL_ROUNDS,
95                PARTIAL_ROUNDS,
96            >>()
97        };
98        debug_assert!(prefix.is_empty(), "Alignment should match");
99        debug_assert!(suffix.is_empty(), "Alignment should match");
100        debug_assert_eq!(shorts.len(), 1);
101        &shorts[0]
102    }
103}
104
105impl<
106    T,
107    const WIDTH: usize,
108    const SBOX_DEGREE: u64,
109    const SBOX_REGISTERS: usize,
110    const HALF_FULL_ROUNDS: usize,
111    const PARTIAL_ROUNDS: usize,
112> BorrowMut<Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>>
113    for [T]
114{
115    fn borrow_mut(
116        &mut self,
117    ) -> &mut Poseidon2Cols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>
118    {
119        // debug_assert_eq!(self.len(), NUM_COLS);
120        let (prefix, shorts, suffix) = unsafe {
121            self.align_to_mut::<Poseidon2Cols<
122                T,
123                WIDTH,
124                SBOX_DEGREE,
125                SBOX_REGISTERS,
126                HALF_FULL_ROUNDS,
127                PARTIAL_ROUNDS,
128            >>()
129        };
130        debug_assert!(prefix.is_empty(), "Alignment should match");
131        debug_assert!(suffix.is_empty(), "Alignment should match");
132        debug_assert_eq!(shorts.len(), 1);
133        &mut shorts[0]
134    }
135}
p3_poseidon2_air/columns.rs

p3_poseidon2_air/
columns.rs
