openvm_poseidon2_air/
lib.rs

1//! This is a wrapper around the Plonky3 [p3_poseidon2_air] used only for integration convenience to
2//! get around some complications with field-specific generics associated with Poseidon2.
3//! Currently it is only intended for use in OpenVM with BabyBear.
4//!
5//! We do not recommend external use of this crate, and suggest using the [p3_poseidon2_air] crate
6//! directly.
7
8use std::sync::Arc;
9
10use openvm_stark_backend::{
11    p3_field::{Field, PrimeField},
12    p3_matrix::dense::RowMajorMatrix,
13};
14pub use openvm_stark_sdk::p3_baby_bear;
15pub use p3_poseidon2;
16use p3_poseidon2::{ExternalLayerConstants, Poseidon2};
17use p3_poseidon2_air::generate_trace_rows;
18pub use p3_poseidon2_air::{self, Poseidon2Air};
19pub use p3_symmetric::{self, Permutation};
20
21mod air;
22mod babybear;
23mod config;
24mod permute;
25
26pub use air::*;
27pub use babybear::*;
28pub use config::*;
29pub use permute::*;
30
31#[cfg(all(feature = "cuda", test))]
32mod cuda_abi;
33
34#[cfg(test)]
35mod tests;
36
37pub const POSEIDON2_WIDTH: usize = 16;
38// NOTE: these constants are for BabyBear only.
39pub const BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS: usize = 4;
40pub const BABY_BEAR_POSEIDON2_FULL_ROUNDS: usize = 8;
41pub const BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS: usize = 13;
42
43// Currently we only support SBOX_DEGREE = 7
44pub const BABY_BEAR_POSEIDON2_SBOX_DEGREE: u64 = 7;
45
46/// `SBOX_REGISTERS` affects the max constraint degree of the AIR. See [p3_poseidon2_air] for more
47/// details.
48#[derive(Debug, Clone)]
49pub struct Poseidon2SubChip<F: Field, const SBOX_REGISTERS: usize> {
50    // This is Arc purely because Poseidon2Air cannot derive Clone
51    pub air: Arc<Poseidon2SubAir<F, SBOX_REGISTERS>>,
52    pub(crate) executor: Poseidon2Executor<F>,
53    pub(crate) constants: Plonky3RoundConstants<F>,
54}
55
56impl<F: PrimeField, const SBOX_REGISTERS: usize> Poseidon2SubChip<F, SBOX_REGISTERS> {
57    pub fn new(constants: Poseidon2Constants<F>) -> Self {
58        let (external_constants, internal_constants) = constants.to_external_internal_constants();
59        Self {
60            air: Arc::new(Poseidon2SubAir::new(constants.into())),
61            executor: Poseidon2Executor::new(external_constants, internal_constants),
62            constants: constants.into(),
63        }
64    }
65
66    pub fn permute(&self, input_state: [F; POSEIDON2_WIDTH]) -> [F; POSEIDON2_WIDTH] {
67        match &self.executor {
68            Poseidon2Executor::BabyBearMds(permuter) => permuter.permute(input_state),
69        }
70    }
71
72    pub fn permute_mut(&self, input_state: &mut [F; POSEIDON2_WIDTH]) {
73        match &self.executor {
74            Poseidon2Executor::BabyBearMds(permuter) => permuter.permute_mut(input_state),
75        };
76    }
77
78    pub fn generate_trace(&self, inputs: Vec<[F; POSEIDON2_WIDTH]>) -> RowMajorMatrix<F>
79    where
80        F: PrimeField,
81    {
82        match self.air.as_ref() {
83            Poseidon2SubAir::BabyBearMds(_) => generate_trace_rows::<
84                F,
85                BabyBearPoseidon2LinearLayers,
86                POSEIDON2_WIDTH,
87                BABY_BEAR_POSEIDON2_SBOX_DEGREE,
88                SBOX_REGISTERS,
89                BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
90                BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
91            >(inputs, &self.constants),
92        }
93    }
94}
95
96#[derive(Clone, Debug)]
97pub enum Poseidon2Executor<F: Field> {
98    BabyBearMds(Plonky3Poseidon2Executor<F, BabyBearPoseidon2LinearLayers>),
99}
100
101impl<F: PrimeField> Poseidon2Executor<F> {
102    pub fn new(
103        external_constants: ExternalLayerConstants<F, POSEIDON2_WIDTH>,
104        internal_constants: Vec<F>,
105    ) -> Self {
106        Self::BabyBearMds(Plonky3Poseidon2Executor::new(
107            external_constants,
108            internal_constants,
109        ))
110    }
111}
112
113pub type Plonky3Poseidon2Executor<F, LinearLayers> = Poseidon2<
114    <F as Field>::Packing,
115    Poseidon2ExternalLayer<F, LinearLayers, POSEIDON2_WIDTH>,
116    Poseidon2InternalLayer<F, LinearLayers>,
117    POSEIDON2_WIDTH,
118    BABY_BEAR_POSEIDON2_SBOX_DEGREE,
119>;
openvm_poseidon2_air/lib.rs

openvm_poseidon2_air/
lib.rs
