openvm_poseidon2_air/
config.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
use openvm_stark_backend::p3_field::{Field, PrimeField32};
use openvm_stark_sdk::p3_baby_bear::BabyBear;
use p3_poseidon2::ExternalLayerConstants;
use p3_poseidon2_air::RoundConstants;

use super::{
    BABYBEAR_BEGIN_EXT_CONSTS, BABYBEAR_END_EXT_CONSTS, BABYBEAR_PARTIAL_CONSTS,
    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS, BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS, POSEIDON2_WIDTH,
};

// Currently only contains round constants, but this struct may contain other configuration parameters in the future.
#[derive(Clone, Copy, Debug)]
pub struct Poseidon2Config<F> {
    pub constants: Poseidon2Constants<F>,
}

impl<F: PrimeField32> Default for Poseidon2Config<F> {
    fn default() -> Self {
        Self {
            constants: default_baby_bear_rc(),
        }
    }
}

#[derive(Clone, Copy, Debug)]
pub struct Poseidon2Constants<F> {
    pub beginning_full_round_constants:
        [[F; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS],
    pub partial_round_constants: [F; BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS],
    pub ending_full_round_constants: [[F; POSEIDON2_WIDTH]; BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS],
}

impl<F: Field> Poseidon2Constants<F> {
    // FIXME[stephenh]: PR https://github.com/Plonky3/Plonky3/pull/588 is now merged.
    // We can remove Poseidon2Constants and use Plonky3RoundConstants directly after updating the plonky3 commit.
    pub fn to_round_constants(&self) -> Plonky3RoundConstants<F> {
        unsafe { std::mem::transmute_copy(self) }
    }

    pub fn to_external_internal_constants(
        &self,
    ) -> (ExternalLayerConstants<F, POSEIDON2_WIDTH>, Vec<F>) {
        (
            ExternalLayerConstants::new(
                self.beginning_full_round_constants.to_vec(),
                self.ending_full_round_constants.to_vec(),
            ),
            self.partial_round_constants.to_vec(),
        )
    }
}

// Round constants for only BabyBear, but we convert to `F` due to some annoyances with generics.
// This should only be used concretely when `F = BabyBear`.
fn default_baby_bear_rc<F: Field>() -> Poseidon2Constants<F> {
    let convert_field = |f: BabyBear| F::from_canonical_u32(f.as_canonical_u32());
    Poseidon2Constants {
        beginning_full_round_constants: BABYBEAR_BEGIN_EXT_CONSTS.map(|x| x.map(convert_field)),
        partial_round_constants: BABYBEAR_PARTIAL_CONSTS.map(convert_field),
        ending_full_round_constants: BABYBEAR_END_EXT_CONSTS.map(|x| x.map(convert_field)),
    }
}

pub type Plonky3RoundConstants<F> = RoundConstants<
    F,
    POSEIDON2_WIDTH,
    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
    BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
>;