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