openvm_poseidon2_air/
permute.rs

1use std::marker::PhantomData;
2
3use derivative::Derivative;
4use openvm_stark_backend::p3_field::{Field, InjectiveMonomial, PrimeCharacteristicRing};
5use p3_poseidon2::{
6    add_rc_and_sbox_generic, mds_light_permutation, ExternalLayer, ExternalLayerConstants,
7    ExternalLayerConstructor, GenericPoseidon2LinearLayers, InternalLayer,
8    InternalLayerConstructor, MDSMat4,
9};
10
11use super::{
12    babybear_internal_linear_layer, BABY_BEAR_POSEIDON2_SBOX_DEGREE, INTERNAL_DIAG_MONTY_16,
13};
14
15const WIDTH: usize = crate::POSEIDON2_WIDTH;
16
17/// Linear layers for BabyBear Poseidon2 using the Plonky3 interfaces, but with a
18/// hand-rolled internal layer to preserve the previous constraint DAG.
19#[derive(Debug, Clone)]
20pub struct BabyBearPoseidon2LinearLayers;
21
22impl GenericPoseidon2LinearLayers<16> for BabyBearPoseidon2LinearLayers {
23    fn internal_linear_layer<R: PrimeCharacteristicRing>(state: &mut [R; WIDTH]) {
24        // Use the old-style internal layer (sum + diag * state[i]) to keep the
25        // Poseidon2 AIR identical to the previous release.
26        babybear_internal_linear_layer(state, &INTERNAL_DIAG_MONTY_16);
27    }
28
29    fn external_linear_layer<R: PrimeCharacteristicRing>(state: &mut [R; WIDTH]) {
30        mds_light_permutation(state, &MDSMat4);
31    }
32}
33
34#[derive(Debug, Derivative)]
35#[derivative(Clone)]
36pub struct Poseidon2InternalLayer<F: PrimeCharacteristicRing, LinearLayers> {
37    pub internal_constants: Vec<F>,
38    _marker: PhantomData<LinearLayers>,
39}
40
41impl<F: Field + PrimeCharacteristicRing, LinearLayers> InternalLayerConstructor<F>
42    for Poseidon2InternalLayer<F, LinearLayers>
43{
44    fn new_from_constants(internal_constants: Vec<F>) -> Self {
45        Self {
46            internal_constants,
47            _marker: PhantomData,
48        }
49    }
50}
51
52impl<
53        F: Field + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>,
54        LinearLayers,
55        const W: usize,
56    > InternalLayer<F, W, BABY_BEAR_POSEIDON2_SBOX_DEGREE>
57    for Poseidon2InternalLayer<F, LinearLayers>
58where
59    LinearLayers: GenericPoseidon2LinearLayers<W>,
60{
61    /// Perform the internal layers of the Poseidon2 permutation on the given state.
62    fn permute_state(&self, state: &mut [F; W]) {
63        self.internal_constants.iter().for_each(|&rc| {
64            add_rc_and_sbox_generic::<_, _, BABY_BEAR_POSEIDON2_SBOX_DEGREE>(&mut state[0], rc);
65            LinearLayers::internal_linear_layer(state);
66        });
67    }
68}
69
70#[derive(Debug, Derivative)]
71#[derivative(Clone)]
72pub struct Poseidon2ExternalLayer<F: PrimeCharacteristicRing, LinearLayers, const W: usize> {
73    pub constants: ExternalLayerConstants<F, W>,
74    _marker: PhantomData<LinearLayers>,
75}
76
77impl<F: Field + PrimeCharacteristicRing, LinearLayers, const W: usize>
78    ExternalLayerConstructor<F, W> for Poseidon2ExternalLayer<F, LinearLayers, W>
79{
80    fn new_from_constants(external_layer_constants: ExternalLayerConstants<F, W>) -> Self {
81        Self {
82            constants: external_layer_constants,
83            _marker: PhantomData,
84        }
85    }
86}
87
88impl<
89        F: Field + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>,
90        LinearLayers,
91        const W: usize,
92    > ExternalLayer<F, W, BABY_BEAR_POSEIDON2_SBOX_DEGREE>
93    for Poseidon2ExternalLayer<F, LinearLayers, W>
94where
95    LinearLayers: GenericPoseidon2LinearLayers<W>,
96{
97    fn permute_state_initial(&self, state: &mut [F; W]) {
98        LinearLayers::external_linear_layer(state);
99        external_permute_state::<F, LinearLayers, W>(state, self.constants.get_initial_constants());
100    }
101
102    fn permute_state_terminal(&self, state: &mut [F; W]) {
103        external_permute_state::<F, LinearLayers, W>(
104            state,
105            self.constants.get_terminal_constants(),
106        );
107    }
108}
109
110fn external_permute_state<
111    F: Field + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>,
112    LinearLayers,
113    const W: usize,
114>(
115    state: &mut [F; W],
116    constants: &[[F; W]],
117) where
118    LinearLayers: GenericPoseidon2LinearLayers<W>,
119{
120    for elem in constants.iter() {
121        state.iter_mut().zip(elem.iter()).for_each(|(s, &rc)| {
122            add_rc_and_sbox_generic::<_, _, BABY_BEAR_POSEIDON2_SBOX_DEGREE>(s, rc)
123        });
124        LinearLayers::external_linear_layer(state);
125    }
126}
openvm_poseidon2_air/permute.rs

openvm_poseidon2_air/
permute.rs
