p3_poseidon2_air/
air.rs

1use core::borrow::Borrow;
2use core::marker::PhantomData;
3
4use p3_air::{Air, AirBuilder, BaseAir};
5use p3_field::{PrimeCharacteristicRing, PrimeField};
6use p3_matrix::Matrix;
7use p3_matrix::dense::RowMajorMatrix;
8use p3_poseidon2::GenericPoseidon2LinearLayers;
9use rand::distr::{Distribution, StandardUniform};
10use rand::rngs::SmallRng;
11use rand::{Rng, SeedableRng};
12
13use crate::columns::{Poseidon2Cols, num_cols};
14use crate::constants::RoundConstants;
15use crate::{FullRound, PartialRound, SBox, generate_trace_rows};
16
17/// Assumes the field size is at least 16 bits.
18#[derive(Debug)]
19pub struct Poseidon2Air<
20    F: PrimeCharacteristicRing,
21    LinearLayers,
22    const WIDTH: usize,
23    const SBOX_DEGREE: u64,
24    const SBOX_REGISTERS: usize,
25    const HALF_FULL_ROUNDS: usize,
26    const PARTIAL_ROUNDS: usize,
27> {
28    pub(crate) constants: RoundConstants<F, WIDTH, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
29    _phantom: PhantomData<LinearLayers>,
30}
31
32impl<
33    F: PrimeCharacteristicRing,
34    LinearLayers,
35    const WIDTH: usize,
36    const SBOX_DEGREE: u64,
37    const SBOX_REGISTERS: usize,
38    const HALF_FULL_ROUNDS: usize,
39    const PARTIAL_ROUNDS: usize,
40> Clone
41    for Poseidon2Air<
42        F,
43        LinearLayers,
44        WIDTH,
45        SBOX_DEGREE,
46        SBOX_REGISTERS,
47        HALF_FULL_ROUNDS,
48        PARTIAL_ROUNDS,
49    >
50{
51    fn clone(&self) -> Self {
52        Self {
53            constants: self.constants.clone(),
54            _phantom: PhantomData,
55        }
56    }
57}
58
59impl<
60    F: PrimeCharacteristicRing,
61    LinearLayers,
62    const WIDTH: usize,
63    const SBOX_DEGREE: u64,
64    const SBOX_REGISTERS: usize,
65    const HALF_FULL_ROUNDS: usize,
66    const PARTIAL_ROUNDS: usize,
67>
68    Poseidon2Air<
69        F,
70        LinearLayers,
71        WIDTH,
72        SBOX_DEGREE,
73        SBOX_REGISTERS,
74        HALF_FULL_ROUNDS,
75        PARTIAL_ROUNDS,
76    >
77{
78    pub const fn new(
79        constants: RoundConstants<F, WIDTH, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>,
80    ) -> Self {
81        Self {
82            constants,
83            _phantom: PhantomData,
84        }
85    }
86
87    pub fn generate_trace_rows(
88        &self,
89        num_hashes: usize,
90        extra_capacity_bits: usize,
91    ) -> RowMajorMatrix<F>
92    where
93        F: PrimeField,
94        LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
95        StandardUniform: Distribution<[F; WIDTH]>,
96    {
97        let mut rng = SmallRng::seed_from_u64(1);
98        let inputs = (0..num_hashes).map(|_| rng.random()).collect();
99        generate_trace_rows::<
100            _,
101            LinearLayers,
102            WIDTH,
103            SBOX_DEGREE,
104            SBOX_REGISTERS,
105            HALF_FULL_ROUNDS,
106            PARTIAL_ROUNDS,
107        >(inputs, &self.constants, extra_capacity_bits)
108    }
109}
110
111impl<
112    F: PrimeCharacteristicRing + Sync,
113    LinearLayers: Sync,
114    const WIDTH: usize,
115    const SBOX_DEGREE: u64,
116    const SBOX_REGISTERS: usize,
117    const HALF_FULL_ROUNDS: usize,
118    const PARTIAL_ROUNDS: usize,
119> BaseAir<F>
120    for Poseidon2Air<
121        F,
122        LinearLayers,
123        WIDTH,
124        SBOX_DEGREE,
125        SBOX_REGISTERS,
126        HALF_FULL_ROUNDS,
127        PARTIAL_ROUNDS,
128    >
129{
130    fn width(&self) -> usize {
131        num_cols::<WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>()
132    }
133}
134
135pub(crate) fn eval<
136    AB: AirBuilder,
137    LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
138    const WIDTH: usize,
139    const SBOX_DEGREE: u64,
140    const SBOX_REGISTERS: usize,
141    const HALF_FULL_ROUNDS: usize,
142    const PARTIAL_ROUNDS: usize,
143>(
144    air: &Poseidon2Air<
145        AB::F,
146        LinearLayers,
147        WIDTH,
148        SBOX_DEGREE,
149        SBOX_REGISTERS,
150        HALF_FULL_ROUNDS,
151        PARTIAL_ROUNDS,
152    >,
153    builder: &mut AB,
154    local: &Poseidon2Cols<
155        AB::Var,
156        WIDTH,
157        SBOX_DEGREE,
158        SBOX_REGISTERS,
159        HALF_FULL_ROUNDS,
160        PARTIAL_ROUNDS,
161    >,
162) {
163    let mut state: [_; WIDTH] = local.inputs.clone().map(|x| x.into());
164
165    LinearLayers::external_linear_layer(&mut state);
166
167    for round in 0..HALF_FULL_ROUNDS {
168        eval_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>(
169            &mut state,
170            &local.beginning_full_rounds[round],
171            &air.constants.beginning_full_round_constants[round],
172            builder,
173        );
174    }
175
176    for round in 0..PARTIAL_ROUNDS {
177        eval_partial_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>(
178            &mut state,
179            &local.partial_rounds[round],
180            &air.constants.partial_round_constants[round],
181            builder,
182        );
183    }
184
185    for round in 0..HALF_FULL_ROUNDS {
186        eval_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>(
187            &mut state,
188            &local.ending_full_rounds[round],
189            &air.constants.ending_full_round_constants[round],
190            builder,
191        );
192    }
193}
194
195impl<
196    AB: AirBuilder,
197    LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
198    const WIDTH: usize,
199    const SBOX_DEGREE: u64,
200    const SBOX_REGISTERS: usize,
201    const HALF_FULL_ROUNDS: usize,
202    const PARTIAL_ROUNDS: usize,
203> Air<AB>
204    for Poseidon2Air<
205        AB::F,
206        LinearLayers,
207        WIDTH,
208        SBOX_DEGREE,
209        SBOX_REGISTERS,
210        HALF_FULL_ROUNDS,
211        PARTIAL_ROUNDS,
212    >
213{
214    #[inline]
215    fn eval(&self, builder: &mut AB) {
216        let main = builder.main();
217        let local = main.row_slice(0).expect("The matrix is empty?");
218        let local = (*local).borrow();
219
220        eval::<_, _, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>(
221            self, builder, local,
222        );
223    }
224}
225
226#[inline]
227fn eval_full_round<
228    AB: AirBuilder,
229    LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
230    const WIDTH: usize,
231    const SBOX_DEGREE: u64,
232    const SBOX_REGISTERS: usize,
233>(
234    state: &mut [AB::Expr; WIDTH],
235    full_round: &FullRound<AB::Var, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>,
236    round_constants: &[AB::F; WIDTH],
237    builder: &mut AB,
238) {
239    for (i, (s, r)) in state.iter_mut().zip(round_constants.iter()).enumerate() {
240        *s += r.clone();
241        eval_sbox(&full_round.sbox[i], s, builder);
242    }
243    LinearLayers::external_linear_layer(state);
244    for (state_i, post_i) in state.iter_mut().zip(&full_round.post) {
245        builder.assert_eq(state_i.clone(), post_i.clone());
246        *state_i = post_i.clone().into();
247    }
248}
249
250#[inline]
251fn eval_partial_round<
252    AB: AirBuilder,
253    LinearLayers: GenericPoseidon2LinearLayers<WIDTH>,
254    const WIDTH: usize,
255    const SBOX_DEGREE: u64,
256    const SBOX_REGISTERS: usize,
257>(
258    state: &mut [AB::Expr; WIDTH],
259    partial_round: &PartialRound<AB::Var, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>,
260    round_constant: &AB::F,
261    builder: &mut AB,
262) {
263    state[0] += round_constant.clone();
264    eval_sbox(&partial_round.sbox, &mut state[0], builder);
265
266    builder.assert_eq(state[0].clone(), partial_round.post_sbox.clone());
267    state[0] = partial_round.post_sbox.clone().into();
268
269    LinearLayers::internal_linear_layer(state);
270}
271
272/// Evaluates the S-box over a degree-1 expression `x`.
273///
274/// # Panics
275///
276/// This method panics if the number of `REGISTERS` is not chosen optimally for the given
277/// `DEGREE` or if the `DEGREE` is not supported by the S-box. The supported degrees are
278/// `3`, `5`, `7`, and `11`.
279#[inline]
280fn eval_sbox<AB, const DEGREE: u64, const REGISTERS: usize>(
281    sbox: &SBox<AB::Var, DEGREE, REGISTERS>,
282    x: &mut AB::Expr,
283    builder: &mut AB,
284) where
285    AB: AirBuilder,
286{
287    *x = match (DEGREE, REGISTERS) {
288        (3, 0) => x.cube(),
289        (5, 0) => x.exp_const_u64::<5>(),
290        (7, 0) => x.exp_const_u64::<7>(),
291        (5, 1) => {
292            let committed_x3 = sbox.0[0].clone().into();
293            let x2 = x.square();
294            builder.assert_eq(committed_x3.clone(), x2.clone() * x.clone());
295            committed_x3 * x2
296        }
297        (7, 1) => {
298            let committed_x3 = sbox.0[0].clone().into();
299            builder.assert_eq(committed_x3.clone(), x.cube());
300            committed_x3.square() * x.clone()
301        }
302        (11, 2) => {
303            let committed_x3 = sbox.0[0].clone().into();
304            let committed_x9 = sbox.0[1].clone().into();
305            let x2 = x.square();
306            builder.assert_eq(committed_x3.clone(), x2.clone() * x.clone());
307            builder.assert_eq(committed_x9.clone(), committed_x3.cube());
308            committed_x9 * x2
309        }
310        _ => panic!(
311            "Unexpected (DEGREE, REGISTERS) of ({}, {})",
312            DEGREE, REGISTERS
313        ),
314    }
315}
p3_poseidon2_air/air.rs

p3_poseidon2_air/
air.rs
