openvm_poseidon2_air/
lib.rs

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

openvm_poseidon2_air/
lib.rs
