openvm_poseidon2_air/
config.rs

1use openvm_stark_backend::p3_field::{Field, PrimeField32};
2use openvm_stark_sdk::p3_baby_bear::BabyBear;
3use p3_poseidon2::ExternalLayerConstants;
4use p3_poseidon2_air::RoundConstants;
5
6use super::{
7    BABYBEAR_BEGIN_EXT_CONSTS, BABYBEAR_END_EXT_CONSTS, BABYBEAR_PARTIAL_CONSTS,
8    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS, BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS, POSEIDON2_WIDTH,
9};
10
11// Currently only contains round constants, but this struct may contain other configuration parameters in the future.
12#[derive(Clone, Copy, Debug)]
13pub struct Poseidon2Config<F> {
14    pub constants: Poseidon2Constants<F>,
15}
16
17impl<F: PrimeField32> Default for Poseidon2Config<F> {
18    fn default() -> Self {
19        Self {
20            constants: default_baby_bear_rc(),
21        }
22    }
23}
24
25#[derive(Clone, Copy, Debug)]
26pub struct Poseidon2Constants<F> {
27    pub beginning_full_round_constants:
28        [[F; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS],
29    pub partial_round_constants: [F; BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS],
30    pub ending_full_round_constants: [[F; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS],
31}
32
33impl<F: Field> From<Poseidon2Constants<F>> for Plonky3RoundConstants<F> {
34    fn from(constants: Poseidon2Constants<F>) -> Self {
35        Plonky3RoundConstants::new(
36            constants.beginning_full_round_constants,
37            constants.partial_round_constants,
38            constants.ending_full_round_constants,
39        )
40    }
41}
42
43impl<F: Field> Poseidon2Constants<F> {
44    pub fn to_external_internal_constants(
45        &self,
46    ) -> (ExternalLayerConstants<F, POSEIDON2_WIDTH>, Vec<F>) {
47        (
48            ExternalLayerConstants::new(
49                self.beginning_full_round_constants.to_vec(),
50                self.ending_full_round_constants.to_vec(),
51            ),
52            self.partial_round_constants.to_vec(),
53        )
54    }
55}
56
57// Round constants for only BabyBear, but we convert to `F` due to some annoyances with generics.
58// This should only be used concretely when `F = BabyBear`.
59fn default_baby_bear_rc<F: Field>() -> Poseidon2Constants<F> {
60    let convert_field = |f: BabyBear| F::from_canonical_u32(f.as_canonical_u32());
61    Poseidon2Constants {
62        beginning_full_round_constants: BABYBEAR_BEGIN_EXT_CONSTS.map(|x| x.map(convert_field)),
63        partial_round_constants: BABYBEAR_PARTIAL_CONSTS.map(convert_field),
64        ending_full_round_constants: BABYBEAR_END_EXT_CONSTS.map(|x| x.map(convert_field)),
65    }
66}
67
68pub type Plonky3RoundConstants<F> = RoundConstants<
69    F,
70    POSEIDON2_WIDTH,
71    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
72    BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
73>;