openvm_circuit/system/poseidon2/
air.rs

1use std::{array::from_fn, borrow::Borrow, sync::Arc};
2
3use derive_new::new;
4use openvm_circuit_primitives::ColumnsAir;
5use openvm_poseidon2_air::{
6    Poseidon2SubAir, BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS, POSEIDON2_WIDTH,
7};
8use openvm_stark_backend::{
9    air_builders::sub::SubAirBuilder,
10    interaction::{InteractionBuilder, LookupBus},
11    p3_air::{Air, BaseAir},
12    p3_field::Field,
13    p3_matrix::Matrix,
14    BaseAirWithPublicValues, PartitionedBaseAir,
15};
16
17use super::columns::Poseidon2PeripheryCols;
18
19/// Poseidon2 Air, VM version.
20///
21/// Carries the subair for subtrace generation. Sticking to the conventions, this struct carries no
22/// state. `direct` determines whether direct interactions are enabled. By default they are on.
23#[derive(Clone, new, Debug)]
24pub struct Poseidon2PeripheryAir<F: Field, const SBOX_REGISTERS: usize> {
25    pub(super) subair: Arc<Poseidon2SubAir<F, SBOX_REGISTERS>>,
26    pub bus: LookupBus,
27}
28
29impl<F: Field, const SBOX_REGISTERS: usize> BaseAirWithPublicValues<F>
30    for Poseidon2PeripheryAir<F, SBOX_REGISTERS>
31{
32}
33impl<F: Field, const SBOX_REGISTERS: usize> PartitionedBaseAir<F>
34    for Poseidon2PeripheryAir<F, SBOX_REGISTERS>
35{
36}
37// No columns provided: columns embed external `Poseidon2SubCols` which doesn't derive
38// `StructReflection`.
39impl<F: Field, const SBOX_REGISTERS: usize> ColumnsAir
40    for Poseidon2PeripheryAir<F, SBOX_REGISTERS>
41{
42}
43impl<F: Field, const SBOX_REGISTERS: usize> BaseAir<F>
44    for Poseidon2PeripheryAir<F, SBOX_REGISTERS>
45{
46    fn width(&self) -> usize {
47        Poseidon2PeripheryCols::<F, SBOX_REGISTERS>::width()
48    }
49}
50
51impl<AB: InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
52    for Poseidon2PeripheryAir<AB::F, SBOX_REGISTERS>
53{
54    fn eval(&self, builder: &mut AB) {
55        let mut sub_builder =
56            SubAirBuilder::<AB, Poseidon2SubAir<AB::F, SBOX_REGISTERS>, AB::F>::new(
57                builder,
58                0..self.subair.width(),
59            );
60        self.subair.eval(&mut sub_builder);
61
62        let main = builder.main();
63        let local = main.row_slice(0).expect("row 0 present");
64        let cols: &Poseidon2PeripheryCols<AB::Var, SBOX_REGISTERS> = (*local).borrow();
65
66        let input: [AB::Var; POSEIDON2_WIDTH] = cols.inner.inputs;
67        let output: [AB::Var; POSEIDON2_WIDTH] =
68            cols.inner.ending_full_rounds[BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1].post;
69        let fields: [_; POSEIDON2_WIDTH + POSEIDON2_WIDTH / 2] = from_fn(|i| {
70            if i < POSEIDON2_WIDTH {
71                input[i]
72            } else {
73                output[i - POSEIDON2_WIDTH]
74            }
75        });
76        self.bus.add_key_with_lookups(builder, fields, cols.mult);
77    }
78}
openvm_circuit/system/poseidon2/air.rs

openvm_circuit/system/poseidon2/
air.rs
