openvm_recursion_circuit/transcript/
poseidon2.rs

1use core::borrow::Borrow;
2use std::{array::from_fn, sync::Arc};
3
4use openvm_circuit_primitives::ColumnsAir;
5use openvm_poseidon2_air::{
6    Poseidon2SubAir, Poseidon2SubCols, BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS,
7};
8use openvm_recursion_circuit_derive::AlignedBorrow;
9use openvm_stark_backend::{
10    air_builders::sub::SubAirBuilder, interaction::InteractionBuilder, BaseAirWithPublicValues,
11    PartitionedBaseAir,
12};
13use p3_air::{Air, AirBuilder, BaseAir};
14use p3_field::Field;
15use p3_matrix::Matrix;
16
17use crate::bus::{
18    Poseidon2CompressBus, Poseidon2CompressMessage, Poseidon2PermuteBus, Poseidon2PermuteMessage,
19};
20
21pub const CHUNK: usize = 8;
22pub use openvm_poseidon2_air::POSEIDON2_WIDTH;
23
24#[repr(C)]
25#[derive(AlignedBorrow)]
26pub struct Poseidon2Cols<T, const SBOX_REGISTERS: usize> {
27    pub inner: Poseidon2SubCols<T, SBOX_REGISTERS>,
28    pub permute_mult: T,
29    pub compress_mult: T,
30}
31
32pub struct Poseidon2Air<F: Field, const SBOX_REGISTERS: usize> {
33    pub subair: Arc<Poseidon2SubAir<F, SBOX_REGISTERS>>,
34    pub poseidon2_permute_bus: Poseidon2PermuteBus,
35    pub poseidon2_compress_bus: Poseidon2CompressBus,
36}
37
38// No columns provided: `Poseidon2Cols` embeds external `Poseidon2SubCols` which doesn't derive
39// `StructReflection`.
40impl<F: Field, const SBOX_REGISTERS: usize> ColumnsAir for Poseidon2Air<F, SBOX_REGISTERS> {}
41
42impl<F: Field, const SBOX_REGISTERS: usize> BaseAir<F> for Poseidon2Air<F, SBOX_REGISTERS> {
43    fn width(&self) -> usize {
44        Poseidon2Cols::<F, SBOX_REGISTERS>::width()
45    }
46}
47
48impl<F: Field, const SBOX_REGISTERS: usize> BaseAirWithPublicValues<F>
49    for Poseidon2Air<F, SBOX_REGISTERS>
50{
51}
52impl<F: Field, const SBOX_REGISTERS: usize> PartitionedBaseAir<F>
53    for Poseidon2Air<F, SBOX_REGISTERS>
54{
55}
56
57impl<AB: AirBuilder + InteractionBuilder, const SBOX_REGISTERS: usize> Air<AB>
58    for Poseidon2Air<AB::F, SBOX_REGISTERS>
59{
60    fn eval(&self, builder: &mut AB) {
61        let main = builder.main();
62        let local = main
63            .row_slice(0)
64            .expect("window should have at least one row");
65        let local: &Poseidon2Cols<AB::Var, SBOX_REGISTERS> = (*local).borrow();
66
67        let mut sub_builder =
68            SubAirBuilder::<AB, Poseidon2SubAir<AB::F, SBOX_REGISTERS>, AB::F>::new(
69                builder,
70                0..self.subair.width(),
71            );
72        self.subair.eval(&mut sub_builder);
73
74        self.poseidon2_permute_bus.add_key_with_lookups(
75            builder,
76            Poseidon2PermuteMessage {
77                input: local.inner.inputs,
78                output: local.inner.ending_full_rounds[BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1]
79                    .post,
80            },
81            local.permute_mult,
82        );
83
84        self.poseidon2_compress_bus.add_key_with_lookups(
85            builder,
86            Poseidon2CompressMessage {
87                input: local.inner.inputs,
88                output: from_fn(|i| {
89                    local.inner.ending_full_rounds[BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS - 1].post[i]
90                }),
91            },
92            local.compress_mult,
93        );
94    }
95}
openvm_recursion_circuit/transcript/poseidon2.rs

openvm_recursion_circuit/transcript/
poseidon2.rs
