openvm_poseidon2_air/
air.rs
1use openvm_stark_backend::{
2 p3_air::{Air, AirBuilder, BaseAir},
3 p3_field::Field,
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, const SBOX_REGISTERS: usize> {
35 BabyBearMds(Plonky3Poseidon2Air<F, BabyBearPoseidon2LinearLayers, SBOX_REGISTERS>),
36}
37
38impl<F: Field, const SBOX_REGISTERS: usize> Poseidon2SubAir<F, SBOX_REGISTERS> {
39 pub fn new(constants: Plonky3RoundConstants<F>) -> Self {
40 Self::BabyBearMds(Plonky3Poseidon2Air::new(constants))
41 }
42}
43
44impl<F: Field, const SBOX_REGISTERS: usize> BaseAir<F> for Poseidon2SubAir<F, SBOX_REGISTERS> {
45 fn width(&self) -> usize {
46 match self {
47 Self::BabyBearMds(air) => air.width(),
48 }
49 }
50}
51
52impl<F: Field, const SBOX_REGISTERS: usize> BaseAirWithPublicValues<F>
53 for Poseidon2SubAir<F, SBOX_REGISTERS>
54{
55}
56impl<F: Field, const SBOX_REGISTERS: usize> PartitionedBaseAir<F>
57 for Poseidon2SubAir<F, SBOX_REGISTERS>
58{
59}
60
61impl<AB: AirBuilder, const SBOX_REGISTERS: usize> Air<AB>
62 for Poseidon2SubAir<AB::F, SBOX_REGISTERS>
63{
64 fn eval(&self, builder: &mut AB) {
65 match self {
66 Self::BabyBearMds(air) => air.eval(builder),
67 }
68 }
69}