openvm_poseidon2_air/
air.rs

1use openvm_circuit_primitives::ColumnsAir;
2use openvm_stark_backend::{
3    p3_air::{Air, AirBuilder, BaseAir},
4    p3_field::{Field, PrimeCharacteristicRing},
5    BaseAirWithPublicValues, PartitionedBaseAir,
6};
7use p3_poseidon2_air::{Poseidon2Air, Poseidon2Cols};
8
9use super::{
10    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS, BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
11    BABY_BEAR_POSEIDON2_SBOX_DEGREE, POSEIDON2_WIDTH,
12};
13use crate::{BabyBearPoseidon2LinearLayers, Plonky3RoundConstants};
14
15pub type Poseidon2SubCols<F, const SBOX_REGISTERS: usize> = Poseidon2Cols<
16    F,
17    POSEIDON2_WIDTH,
18    BABY_BEAR_POSEIDON2_SBOX_DEGREE,
19    SBOX_REGISTERS,
20    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
21    BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
22>;
23
24pub type Plonky3Poseidon2Air<F, LinearLayers, const SBOX_REGISTERS: usize> = Poseidon2Air<
25    F,
26    LinearLayers,
27    POSEIDON2_WIDTH,
28    BABY_BEAR_POSEIDON2_SBOX_DEGREE,
29    SBOX_REGISTERS,
30    BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
31    BABY_BEAR_POSEIDON2_PARTIAL_ROUNDS,
32>;
33
34#[derive(Debug)]
35pub enum Poseidon2SubAir<F: Field + PrimeCharacteristicRing, const SBOX_REGISTERS: usize> {
36    BabyBearMds(Plonky3Poseidon2Air<F, BabyBearPoseidon2LinearLayers, SBOX_REGISTERS>),
37}
38
39impl<F: Field + PrimeCharacteristicRing, const SBOX_REGISTERS: usize>
40    Poseidon2SubAir<F, SBOX_REGISTERS>
41{
42    pub fn new(constants: Plonky3RoundConstants<F>) -> Self {
43        Self::BabyBearMds(Plonky3Poseidon2Air::new(constants))
44    }
45}
46
47impl<F: Field + PrimeCharacteristicRing, const SBOX_REGISTERS: usize> BaseAir<F>
48    for Poseidon2SubAir<F, SBOX_REGISTERS>
49{
50    fn width(&self) -> usize {
51        match self {
52            Self::BabyBearMds(air) => air.width(),
53        }
54    }
55}
56
57impl<F: Field + PrimeCharacteristicRing, const SBOX_REGISTERS: usize> BaseAirWithPublicValues<F>
58    for Poseidon2SubAir<F, SBOX_REGISTERS>
59{
60}
61impl<F: Field + PrimeCharacteristicRing, const SBOX_REGISTERS: usize> PartitionedBaseAir<F>
62    for Poseidon2SubAir<F, SBOX_REGISTERS>
63{
64}
65// No columns provided: `Poseidon2Cols` (from external `p3-poseidon2-air`) doesn't derive
66// `StructReflection`.
67impl<F: Field + PrimeCharacteristicRing, const SBOX_REGISTERS: usize> ColumnsAir
68    for Poseidon2SubAir<F, SBOX_REGISTERS>
69{
70}
71
72impl<AB: AirBuilder, const SBOX_REGISTERS: usize> Air<AB> for Poseidon2SubAir<AB::F, SBOX_REGISTERS>
73where
74    AB::F: Field + PrimeCharacteristicRing,
75{
76    fn eval(&self, builder: &mut AB) {
77        match self {
78            Self::BabyBearMds(air) => air.eval(builder),
79        }
80    }
81}
openvm_poseidon2_air/air.rs

openvm_poseidon2_air/
air.rs
