openvm_circuit/system/poseidon2/
air.rs

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

openvm_circuit/system/poseidon2/
air.rs
