openvm_deferral_circuit/poseidon2/
air.rs

1use std::{borrow::Borrow, iter::once, sync::Arc};
2
3use openvm_circuit_primitives::{AlignedBorrow, ColumnsAir};
4use openvm_poseidon2_air::{
5    Poseidon2Config, Poseidon2SubAir, Poseidon2SubCols, BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
6};
7use openvm_stark_backend::{
8    air_builders::sub::SubAirBuilder,
9    interaction::{InteractionBuilder, LookupBus},
10    p3_air::{Air, AirBuilder, BaseAir},
11    p3_matrix::Matrix,
12    BaseAirWithPublicValues, PartitionedBaseAir,
13};
14use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
15use p3_field::{Field, PrimeCharacteristicRing};
16
17use super::SBOX_REGISTERS;
18
19#[repr(C)]
20#[derive(AlignedBorrow)]
21pub struct DeferralPoseidon2Cols<T> {
22    pub inner: Poseidon2SubCols<T, SBOX_REGISTERS>,
23    pub compress_mult: T,
24    pub capacity_mult: T,
25}
26
27pub struct DeferralPoseidon2Air<F: Field> {
28    pub subair: Arc<Poseidon2SubAir<F, SBOX_REGISTERS>>,
29    pub bus: LookupBus,
30}
31
32// No columns provided: `DeferralPoseidon2Cols` embeds external `Poseidon2SubCols` which doesn't
33// derive `StructReflection`.
34impl<F: Field> ColumnsAir for DeferralPoseidon2Air<F> {}
35
36impl<F: Field> DeferralPoseidon2Air<F> {
37    pub fn new(config: Poseidon2Config<F>, bus: LookupBus) -> Self {
38        Self {
39            subair: Arc::new(Poseidon2SubAir::new(config.constants.into())),
40            bus,
41        }
42    }
43}
44
45impl<F: Field> BaseAir<F> for DeferralPoseidon2Air<F> {
46    fn width(&self) -> usize {
47        DeferralPoseidon2Cols::<F>::width()
48    }
49}
50impl<F: Field> BaseAirWithPublicValues<F> for DeferralPoseidon2Air<F> {}
51impl<F: Field> PartitionedBaseAir<F> for DeferralPoseidon2Air<F> {}
52
53impl<AB: AirBuilder + InteractionBuilder> Air<AB> for DeferralPoseidon2Air<AB::F> {
54    fn eval(&self, builder: &mut AB) {
55        let main = builder.main();
56        let local = main
57            .row_slice(0)
58            .expect("window should have at least one row");
59        let local: &DeferralPoseidon2Cols<AB::Var> = (*local).borrow();
60
61        let mut sub_builder =
62            SubAirBuilder::<AB, Poseidon2SubAir<AB::F, SBOX_REGISTERS>, AB::F>::new(
63                builder,
64                0..self.subair.width(),
65            );
66        self.subair.eval(&mut sub_builder);
67
68        let inputs = local.inner.inputs;
69        let compress_res = &local.inner.ending_full_rounds
70            [BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1]
71            .post[..DIGEST_SIZE];
72        let capacity_res = &local.inner.ending_full_rounds
73            [BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1]
74            .post[DIGEST_SIZE..];
75
76        self.bus.add_key_with_lookups(
77            builder,
78            inputs
79                .into_iter()
80                .map(Into::into)
81                .chain(compress_res.iter().copied().map(Into::into))
82                .chain(once(AB::Expr::ONE)),
83            local.compress_mult,
84        );
85
86        self.bus.add_key_with_lookups(
87            builder,
88            inputs
89                .into_iter()
90                .map(Into::into)
91                .chain(capacity_res.iter().copied().map(Into::into))
92                .chain(once(AB::Expr::ZERO)),
93            local.capacity_mult,
94        );
95    }
96}
openvm_deferral_circuit/poseidon2/air.rs

openvm_deferral_circuit/poseidon2/
air.rs
