openvm_poseidon2_air/
permute.rs

1use std::{any::TypeId, marker::PhantomData};
2
3use derivative::Derivative;
4use openvm_stark_backend::p3_field::FieldAlgebra;
5use openvm_stark_sdk::p3_baby_bear::{BabyBear, BabyBearInternalLayerParameters};
6use p3_monty_31::InternalLayerBaseParameters;
7use p3_poseidon2::{
8    add_rc_and_sbox_generic, mds_light_permutation, ExternalLayer, ExternalLayerConstants,
9    ExternalLayerConstructor, GenericPoseidon2LinearLayers, InternalLayer,
10    InternalLayerConstructor, MDSMat4,
11};
12
13use super::{babybear_internal_linear_layer, BABY_BEAR_POSEIDON2_SBOX_DEGREE};
14
15const WIDTH: usize = crate::POSEIDON2_WIDTH;
16
17pub trait Poseidon2MatrixConfig: Clone + Sync {
18    fn int_diag_m1_matrix<F: FieldAlgebra>() -> [F; WIDTH];
19}
20
21/// This type needs to implement GenericPoseidon2LinearLayers generic in F so that our
22/// Poseidon2SubAir can also be generic in F, but in reality each implementation of this struct's
23/// functions should be field specific. To circumvent this, Poseidon2LinearLayers is generic in F
24/// but **currently requires** that F is BabyBear.
25#[derive(Debug, Clone)]
26pub struct BabyBearPoseidon2LinearLayers;
27
28// This is the same as the implementation for
29// GenericPoseidon2LinearLayersMonty31<BabyBearParameters, BabyBearInternalLayerParameters> except
30// that we drop the clause that FA needs be multipliable by BabyBear.
31// TODO[jpw/stephen]: This is clearly not the best way to do this, but it would
32// require some reworking in plonky3 to get around the generics.
33impl<FA: FieldAlgebra> GenericPoseidon2LinearLayers<FA, WIDTH> for BabyBearPoseidon2LinearLayers {
34    fn internal_linear_layer(state: &mut [FA; WIDTH]) {
35        let diag_m1_matrix = &<BabyBearInternalLayerParameters as InternalLayerBaseParameters<
36            _,
37            16,
38        >>::INTERNAL_DIAG_MONTY;
39        assert_eq!(
40            TypeId::of::<FA::F>(),
41            TypeId::of::<BabyBear>(),
42            "BabyBear is the only supported field type"
43        );
44        // SAFETY: TypeId check above ensures FA::F is BabyBear, so transmute is valid
45        let diag_m1_matrix =
46            unsafe { std::mem::transmute::<&[BabyBear; WIDTH], &[FA::F; WIDTH]>(diag_m1_matrix) };
47        babybear_internal_linear_layer(state, diag_m1_matrix);
48    }
49
50    fn external_linear_layer(state: &mut [FA; WIDTH]) {
51        mds_light_permutation(state, &MDSMat4);
52    }
53}
54
55// Below are generic implementations of the Poseidon2 Internal and External Layers
56// generic in the field. These are currently used for the runtime poseidon2
57// execution even though they are less optimized than the Monty31 specific
58// implementations in Plonky3. We could use those more optimized implementations,
59// but it would require many unsafe transmutes.
60
61#[derive(Debug, Derivative)]
62#[derivative(Clone)]
63pub struct Poseidon2InternalLayer<F: FieldAlgebra, LinearLayers> {
64    pub internal_constants: Vec<F>,
65    _marker: PhantomData<LinearLayers>,
66}
67
68impl<AF: FieldAlgebra, LinearLayers> InternalLayerConstructor<AF>
69    for Poseidon2InternalLayer<AF::F, LinearLayers>
70{
71    fn new_from_constants(internal_constants: Vec<AF::F>) -> Self {
72        Self {
73            internal_constants,
74            _marker: PhantomData,
75        }
76    }
77}
78
79impl<FA: FieldAlgebra, LinearLayers, const WIDTH: usize>
80    InternalLayer<FA, WIDTH, BABY_BEAR_POSEIDON2_SBOX_DEGREE>
81    for Poseidon2InternalLayer<FA::F, LinearLayers>
82where
83    LinearLayers: GenericPoseidon2LinearLayers<FA, WIDTH>,
84{
85    /// Perform the internal layers of the Poseidon2 permutation on the given state.
86    fn permute_state(&self, state: &mut [FA; WIDTH]) {
87        self.internal_constants.iter().for_each(|&rc| {
88            add_rc_and_sbox_generic::<_, BABY_BEAR_POSEIDON2_SBOX_DEGREE>(&mut state[0], rc);
89            LinearLayers::internal_linear_layer(state);
90        })
91    }
92}
93
94#[derive(Debug, Derivative)]
95#[derivative(Clone)]
96pub struct Poseidon2ExternalLayer<F: FieldAlgebra, LinearLayers, const WIDTH: usize> {
97    pub constants: ExternalLayerConstants<F, WIDTH>,
98    _marker: PhantomData<LinearLayers>,
99}
100
101impl<FA: FieldAlgebra, LinearLayers, const WIDTH: usize> ExternalLayerConstructor<FA, WIDTH>
102    for Poseidon2ExternalLayer<FA::F, LinearLayers, WIDTH>
103{
104    fn new_from_constants(external_layer_constants: ExternalLayerConstants<FA::F, WIDTH>) -> Self {
105        Self {
106            constants: external_layer_constants,
107            _marker: PhantomData,
108        }
109    }
110}
111
112impl<FA: FieldAlgebra, LinearLayers, const WIDTH: usize>
113    ExternalLayer<FA, WIDTH, BABY_BEAR_POSEIDON2_SBOX_DEGREE>
114    for Poseidon2ExternalLayer<FA::F, LinearLayers, WIDTH>
115where
116    LinearLayers: GenericPoseidon2LinearLayers<FA, WIDTH>,
117{
118    fn permute_state_initial(&self, state: &mut [FA; WIDTH]) {
119        LinearLayers::external_linear_layer(state);
120        external_permute_state::<FA, LinearLayers, WIDTH>(
121            state,
122            self.constants.get_initial_constants(),
123        );
124    }
125
126    fn permute_state_terminal(&self, state: &mut [FA; WIDTH]) {
127        external_permute_state::<FA, LinearLayers, WIDTH>(
128            state,
129            self.constants.get_terminal_constants(),
130        );
131    }
132}
133
134fn external_permute_state<FA: FieldAlgebra, LinearLayers, const WIDTH: usize>(
135    state: &mut [FA; WIDTH],
136    constants: &[[FA::F; WIDTH]],
137) where
138    LinearLayers: GenericPoseidon2LinearLayers<FA, WIDTH>,
139{
140    for elem in constants.iter() {
141        state.iter_mut().zip(elem.iter()).for_each(|(s, &rc)| {
142            add_rc_and_sbox_generic::<_, BABY_BEAR_POSEIDON2_SBOX_DEGREE>(s, rc)
143        });
144        LinearLayers::external_linear_layer(state);
145    }
146}
openvm_poseidon2_air/permute.rs

openvm_poseidon2_air/
permute.rs
