openvm_poseidon2_air/
permute.rs

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

openvm_poseidon2_air/
permute.rs
