openvm_poseidon2_air/
air.rs

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