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::{InjectiveMonomial, 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<
50    F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>,
51    const SBOX_REGISTERS: usize,
52> {
53    // This is Arc purely because Poseidon2Air cannot derive Clone
54    pub air: Arc<Poseidon2SubAir<F, SBOX_REGISTERS>>,
55    pub(crate) executor: Poseidon2Executor<F>,
56    pub(crate) constants: Plonky3RoundConstants<F>,
57}
58
59impl<
60        F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>,
61        const SBOX_REGISTERS: usize,
62    > Poseidon2SubChip<F, SBOX_REGISTERS>
63{
64    pub fn new(constants: Poseidon2Constants<F>) -> Self {
65        let (external_constants, internal_constants) = constants.to_external_internal_constants();
66        Self {
67            air: Arc::new(Poseidon2SubAir::new(constants.into())),
68            executor: Poseidon2Executor::new(external_constants, internal_constants),
69            constants: constants.into(),
70        }
71    }
72
73    pub fn permute(&self, input_state: [F; POSEIDON2_WIDTH]) -> [F; POSEIDON2_WIDTH] {
74        match &self.executor {
75            Poseidon2Executor::BabyBearMds(permuter) => permuter.permute(input_state),
76        }
77    }
78
79    pub fn permute_mut(&self, input_state: &mut [F; POSEIDON2_WIDTH]) {
80        match &self.executor {
81            Poseidon2Executor::BabyBearMds(permuter) => permuter.permute_mut(input_state),
82        };
83    }
84
85    pub fn generate_trace(&self, inputs: Vec<[F; POSEIDON2_WIDTH]>) -> RowMajorMatrix<F>
86    where
87        F: PrimeField,
88    {
89        match self.air.as_ref() {
90            // The third argument `extra_capacity_bits` pre-allocates Vec capacity for
91            // the LDE blowup during PCS commit (capacity = size << extra_capacity_bits).
92            // This optimization only helps if the returned matrix flows directly to
93            // the prover. Currently in OpenVM, this sub-trace is copied into wider
94            // matrices that include additional columns, so the extra capacity would
95            // be discarded.
96            Poseidon2SubAir::BabyBearMds(_) => generate_trace_rows::<
97                F,
98                BabyBearPoseidon2LinearLayers,
99                POSEIDON2_WIDTH,
100                BABY_BEAR_POSEIDON2_SBOX_DEGREE,
101                SBOX_REGISTERS,
102                BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
103                BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
104            >(inputs, &self.constants, 0),
105        }
106    }
107}
108
109#[derive(Clone, Debug)]
110pub enum Poseidon2Executor<F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>> {
111    BabyBearMds(Plonky3Poseidon2Executor<F, BabyBearPoseidon2LinearLayers>),
112}
113
114impl<F: PrimeField + InjectiveMonomial<BABY_BEAR_POSEIDON2_SBOX_DEGREE>> Poseidon2Executor<F> {
115    pub fn new(
116        external_constants: ExternalLayerConstants<F, POSEIDON2_WIDTH>,
117        internal_constants: Vec<F>,
118    ) -> Self {
119        Self::BabyBearMds(Plonky3Poseidon2Executor::new(
120            external_constants,
121            internal_constants,
122        ))
123    }
124}
125
126pub type Plonky3Poseidon2Executor<F, LinearLayers> = Poseidon2<
127    F,
128    Poseidon2ExternalLayer<F, LinearLayers, POSEIDON2_WIDTH>,
129    Poseidon2InternalLayer<F, LinearLayers>,
130    POSEIDON2_WIDTH,
131    BABY_BEAR_POSEIDON2_SBOX_DEGREE,
132>;
openvm_poseidon2_air/lib.rs

openvm_poseidon2_air/
lib.rs
