p3_poseidon2_air/
constants.rs

1use alloc::vec::Vec;
2
3use p3_field::Field;
4use rand::distributions::{Distribution, Standard};
5use rand::Rng;
6
7/// Round constants for Poseidon2, in a format that's convenient for the AIR.
8#[derive(Debug, Clone)]
9pub struct RoundConstants<
10    F: Field,
11    const WIDTH: usize,
12    const HALF_FULL_ROUNDS: usize,
13    const PARTIAL_ROUNDS: usize,
14> {
15    pub(crate) beginning_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS],
16    pub(crate) partial_round_constants: [F; PARTIAL_ROUNDS],
17    pub(crate) ending_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS],
18}
19
20impl<F: Field, const WIDTH: usize, const HALF_FULL_ROUNDS: usize, const PARTIAL_ROUNDS: usize>
21    RoundConstants<F, WIDTH, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>
22{
23    pub fn new(
24        beginning_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS],
25        partial_round_constants: [F; PARTIAL_ROUNDS],
26        ending_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS],
27    ) -> Self {
28        Self {
29            beginning_full_round_constants,
30            partial_round_constants,
31            ending_full_round_constants,
32        }
33    }
34
35    pub fn from_rng<R: Rng>(rng: &mut R) -> Self
36    where
37        Standard: Distribution<F> + Distribution<[F; WIDTH]>,
38    {
39        let beginning_full_round_constants = rng
40            .sample_iter(Standard)
41            .take(HALF_FULL_ROUNDS)
42            .collect::<Vec<[F; WIDTH]>>()
43            .try_into()
44            .unwrap();
45        let partial_round_constants = rng
46            .sample_iter(Standard)
47            .take(PARTIAL_ROUNDS)
48            .collect::<Vec<F>>()
49            .try_into()
50            .unwrap();
51        let ending_full_round_constants = rng
52            .sample_iter(Standard)
53            .take(HALF_FULL_ROUNDS)
54            .collect::<Vec<[F; WIDTH]>>()
55            .try_into()
56            .unwrap();
57        Self {
58            beginning_full_round_constants,
59            partial_round_constants,
60            ending_full_round_constants,
61        }
62    }
63}