From 800737a1bbe2713cd85094390dc121bbdd34cd91 Mon Sep 17 00:00:00 2001 From: Sai Deng Date: Wed, 8 Oct 2025 21:24:47 -0700 Subject: [PATCH 1/9] copy --- poseidon2-air/Cargo.toml | 41 +++++ poseidon2-air/src/air.rs | 288 ++++++++++++++++++++++++++++++++ poseidon2-air/src/columns.rs | 162 ++++++++++++++++++ poseidon2-air/src/constants.rs | 47 ++++++ poseidon2-air/src/generation.rs | 288 ++++++++++++++++++++++++++++++++ poseidon2-air/src/lib.rs | 17 ++ poseidon2-air/src/vectorized.rs | 275 ++++++++++++++++++++++++++++++ 7 files changed, 1118 insertions(+) create mode 100644 poseidon2-air/Cargo.toml create mode 100644 poseidon2-air/src/air.rs create mode 100644 poseidon2-air/src/columns.rs create mode 100644 poseidon2-air/src/constants.rs create mode 100644 poseidon2-air/src/generation.rs create mode 100644 poseidon2-air/src/lib.rs create mode 100644 poseidon2-air/src/vectorized.rs diff --git a/poseidon2-air/Cargo.toml b/poseidon2-air/Cargo.toml new file mode 100644 index 0000000..f782c2c --- /dev/null +++ b/poseidon2-air/Cargo.toml @@ -0,0 +1,41 @@ +[package] +name = "p3-poseidon2-air" +description = "An AIR implementation of the Poseidon2 cryptographic hash function for use in zero-knowledge proof systems." +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +homepage.workspace = true +keywords.workspace = true +categories.workspace = true + +[dependencies] +p3-air.workspace = true +p3-field.workspace = true +p3-matrix.workspace = true +p3-maybe-rayon.workspace = true +p3-poseidon2.workspace = true + +rand.workspace = true +tracing.workspace = true + +[target.'cfg(target_family = "unix")'.dev-dependencies] +tikv-jemallocator = "0.6" + +[dev-dependencies] +p3-baby-bear.workspace = true +p3-challenger.workspace = true +p3-commit.workspace = true +p3-dft.workspace = true +p3-fri.workspace = true +p3-keccak.workspace = true +p3-koala-bear.workspace = true +p3-merkle-tree.workspace = true +p3-symmetric.workspace = true +p3-uni-stark.workspace = true + +tracing-forest = { workspace = true, features = ["ansi", "smallvec"] } +tracing-subscriber = { workspace = true, features = ["std", "env-filter"] } + +[features] +parallel = ["p3-maybe-rayon/parallel"] diff --git a/poseidon2-air/src/air.rs b/poseidon2-air/src/air.rs new file mode 100644 index 0000000..cfa9e5b --- /dev/null +++ b/poseidon2-air/src/air.rs @@ -0,0 +1,288 @@ +use core::borrow::Borrow; +use core::marker::PhantomData; + +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{PrimeCharacteristicRing, PrimeField}; +use p3_matrix::Matrix; +use p3_matrix::dense::RowMajorMatrix; +use p3_poseidon2::GenericPoseidon2LinearLayers; +use rand::distr::{Distribution, StandardUniform}; +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; + +use crate::columns::{Poseidon2Cols, num_cols}; +use crate::constants::RoundConstants; +use crate::{FullRound, PartialRound, SBox, generate_trace_rows}; + +/// Assumes the field size is at least 16 bits. +#[derive(Debug)] +pub struct Poseidon2Air< + F: PrimeCharacteristicRing, + LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> { + pub(crate) constants: RoundConstants, + _phantom: PhantomData, +} + +impl< + F: PrimeCharacteristicRing, + LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> + Poseidon2Air< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > +{ + pub const fn new( + constants: RoundConstants, + ) -> Self { + Self { + constants, + _phantom: PhantomData, + } + } + + pub fn generate_trace_rows( + &self, + num_hashes: usize, + extra_capacity_bits: usize, + ) -> RowMajorMatrix + where + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + StandardUniform: Distribution<[F; WIDTH]>, + { + let mut rng = SmallRng::seed_from_u64(1); + let inputs = (0..num_hashes).map(|_| rng.random()).collect(); + generate_trace_rows::< + _, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(inputs, &self.constants, extra_capacity_bits) + } +} + +impl< + F: PrimeCharacteristicRing + Sync, + LinearLayers: Sync, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> BaseAir + for Poseidon2Air< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > +{ + fn width(&self) -> usize { + num_cols::() + } +} + +pub(crate) fn eval< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>( + air: &Poseidon2Air< + AB::F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + builder: &mut AB, + local: &Poseidon2Cols< + AB::Var, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, +) { + let mut state: [_; WIDTH] = local.inputs.clone().map(|x| x.into()); + + LinearLayers::external_linear_layer(&mut state); + + for round in 0..HALF_FULL_ROUNDS { + eval_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( + &mut state, + &local.beginning_full_rounds[round], + &air.constants.beginning_full_round_constants[round], + builder, + ); + } + + for round in 0..PARTIAL_ROUNDS { + eval_partial_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( + &mut state, + &local.partial_rounds[round], + &air.constants.partial_round_constants[round], + builder, + ); + } + + for round in 0..HALF_FULL_ROUNDS { + eval_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( + &mut state, + &local.ending_full_rounds[round], + &air.constants.ending_full_round_constants[round], + builder, + ); + } +} + +impl< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> Air + for Poseidon2Air< + AB::F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > +{ + #[inline] + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0).expect("The matrix is empty?"); + let local = (*local).borrow(); + + eval::<_, _, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>( + self, builder, local, + ); + } +} + +#[inline] +fn eval_full_round< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, +>( + state: &mut [AB::Expr; WIDTH], + full_round: &FullRound, + round_constants: &[AB::F; WIDTH], + builder: &mut AB, +) { + for (i, (s, r)) in state.iter_mut().zip(round_constants.iter()).enumerate() { + *s += r.clone(); + eval_sbox(&full_round.sbox[i], s, builder); + } + LinearLayers::external_linear_layer(state); + for (state_i, post_i) in state.iter_mut().zip(&full_round.post) { + builder.assert_eq(state_i.clone(), post_i.clone()); + *state_i = post_i.clone().into(); + } +} + +#[inline] +fn eval_partial_round< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, +>( + state: &mut [AB::Expr; WIDTH], + partial_round: &PartialRound, + round_constant: &AB::F, + builder: &mut AB, +) { + state[0] += round_constant.clone(); + eval_sbox(&partial_round.sbox, &mut state[0], builder); + + builder.assert_eq(state[0].clone(), partial_round.post_sbox.clone()); + state[0] = partial_round.post_sbox.clone().into(); + + LinearLayers::internal_linear_layer(state); +} + +/// Evaluates the S-box over a degree-1 expression `x`. +/// +/// # Panics +/// +/// This method panics if the number of `REGISTERS` is not chosen optimally for the given +/// `DEGREE` or if the `DEGREE` is not supported by the S-box. The supported degrees are +/// `3`, `5`, `7`, and `11`. +#[inline] +fn eval_sbox( + sbox: &SBox, + x: &mut AB::Expr, + builder: &mut AB, +) where + AB: AirBuilder, +{ + *x = match (DEGREE, REGISTERS) { + (3, 0) => x.cube(), + (5, 0) => x.exp_const_u64::<5>(), + (7, 0) => x.exp_const_u64::<7>(), + (5, 1) => { + let committed_x3 = sbox.0[0].clone().into(); + let x2 = x.square(); + builder.assert_eq(committed_x3.clone(), x2.clone() * x.clone()); + committed_x3 * x2 + } + (7, 1) => { + let committed_x3 = sbox.0[0].clone().into(); + builder.assert_eq(committed_x3.clone(), x.cube()); + committed_x3.square() * x.clone() + } + (11, 2) => { + let committed_x3 = sbox.0[0].clone().into(); + let committed_x9 = sbox.0[1].clone().into(); + let x2 = x.square(); + builder.assert_eq(committed_x3.clone(), x2.clone() * x.clone()); + builder.assert_eq(committed_x9.clone(), committed_x3.cube()); + committed_x9 * x2 + } + _ => panic!( + "Unexpected (DEGREE, REGISTERS) of ({}, {})", + DEGREE, REGISTERS + ), + } +} diff --git a/poseidon2-air/src/columns.rs b/poseidon2-air/src/columns.rs new file mode 100644 index 0000000..082f318 --- /dev/null +++ b/poseidon2-air/src/columns.rs @@ -0,0 +1,162 @@ +use core::borrow::{Borrow, BorrowMut}; +use core::mem::size_of; + +/// Columns for a Poseidon2 AIR which computes one permutation per row. +/// +/// The columns of the STARK are divided into the three different round sections of the Poseidon2 +/// Permutation: beginning full rounds, partial rounds, and ending full rounds. For the full +/// rounds we store an [`SBox`] columnset for each state variable, and for the partial rounds we +/// store only for the first state variable. Because the matrix multiplications are linear +/// functions, we need only keep auxiliary columns for the S-box computations. +#[repr(C)] +pub struct Poseidon2Cols< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> { + pub export: T, + + pub inputs: [T; WIDTH], + + /// Beginning Full Rounds + pub beginning_full_rounds: [FullRound; HALF_FULL_ROUNDS], + + /// Partial Rounds + pub partial_rounds: [PartialRound; PARTIAL_ROUNDS], + + /// Ending Full Rounds + pub ending_full_rounds: [FullRound; HALF_FULL_ROUNDS], +} + +/// Full round columns. +#[repr(C)] +pub struct FullRound { + /// Possible intermediate results within each S-box. + pub sbox: [SBox; WIDTH], + /// The post-state, i.e. the entire layer after this full round. + pub post: [T; WIDTH], +} + +/// Partial round columns. +#[repr(C)] +pub struct PartialRound +{ + /// Possible intermediate results within the S-box. + pub sbox: SBox, + /// The output of the S-box. + pub post_sbox: T, +} + +/// Possible intermediate results within an S-box. +/// +/// Use this column-set for an S-box that can be computed with `REGISTERS`-many intermediate results +/// (not counting the final output). The S-box is checked to ensure that `REGISTERS` is the optimal +/// number of registers for the given `DEGREE` for the degrees given in the Poseidon2 paper: +/// `3`, `5`, `7`, and `11`. See `eval_sbox` for more information. +#[repr(C)] +pub struct SBox(pub [T; REGISTERS]); + +pub const fn num_cols< + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>() -> usize { + size_of::>( + ) +} + +pub const fn make_col_map< + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>() -> Poseidon2Cols { + todo!() + // let indices_arr = indices_arr::< + // { num_cols::() }, + // >(); + // unsafe { + // transmute::< + // [usize; + // num_cols::()], + // Poseidon2Cols< + // usize, + // WIDTH, + // SBOX_DEGREE, + // SBOX_REGISTERS, + // HALF_FULL_ROUNDS, + // PARTIAL_ROUNDS, + // >, + // >(indices_arr) + // } +} + +impl< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> Borrow> + for [T] +{ + fn borrow( + &self, + ) -> &Poseidon2Cols + { + // debug_assert_eq!(self.len(), NUM_COLS); + let (prefix, shorts, suffix) = unsafe { + self.align_to::>() + }; + debug_assert!(prefix.is_empty(), "Alignment should match"); + debug_assert!(suffix.is_empty(), "Alignment should match"); + debug_assert_eq!(shorts.len(), 1); + &shorts[0] + } +} + +impl< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> BorrowMut> + for [T] +{ + fn borrow_mut( + &mut self, + ) -> &mut Poseidon2Cols + { + // debug_assert_eq!(self.len(), NUM_COLS); + let (prefix, shorts, suffix) = unsafe { + self.align_to_mut::>() + }; + debug_assert!(prefix.is_empty(), "Alignment should match"); + debug_assert!(suffix.is_empty(), "Alignment should match"); + debug_assert_eq!(shorts.len(), 1); + &mut shorts[0] + } +} diff --git a/poseidon2-air/src/constants.rs b/poseidon2-air/src/constants.rs new file mode 100644 index 0000000..5e50607 --- /dev/null +++ b/poseidon2-air/src/constants.rs @@ -0,0 +1,47 @@ +use p3_field::PrimeCharacteristicRing; +use rand::Rng; +use rand::distr::{Distribution, StandardUniform}; + +/// Round constants for Poseidon2, in a format that's convenient for the AIR. +#[derive(Debug, Clone)] +pub struct RoundConstants< + F: PrimeCharacteristicRing, + const WIDTH: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> { + pub(crate) beginning_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], + pub(crate) partial_round_constants: [F; PARTIAL_ROUNDS], + pub(crate) ending_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], +} + +impl< + F: PrimeCharacteristicRing, + const WIDTH: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> RoundConstants +{ + pub const fn new( + beginning_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], + partial_round_constants: [F; PARTIAL_ROUNDS], + ending_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], + ) -> Self { + Self { + beginning_full_round_constants, + partial_round_constants, + ending_full_round_constants, + } + } + + pub fn from_rng(rng: &mut R) -> Self + where + StandardUniform: Distribution + Distribution<[F; WIDTH]>, + { + Self { + beginning_full_round_constants: core::array::from_fn(|_| rng.sample(StandardUniform)), + partial_round_constants: core::array::from_fn(|_| rng.sample(StandardUniform)), + ending_full_round_constants: core::array::from_fn(|_| rng.sample(StandardUniform)), + } + } +} diff --git a/poseidon2-air/src/generation.rs b/poseidon2-air/src/generation.rs new file mode 100644 index 0000000..ebef843 --- /dev/null +++ b/poseidon2-air/src/generation.rs @@ -0,0 +1,288 @@ +use alloc::vec::Vec; +use core::mem::MaybeUninit; + +use p3_field::PrimeField; +use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixViewMut}; +use p3_maybe_rayon::prelude::*; +use p3_poseidon2::GenericPoseidon2LinearLayers; +use tracing::instrument; + +use crate::columns::{Poseidon2Cols, num_cols}; +use crate::{FullRound, PartialRound, RoundConstants, SBox}; + +#[instrument(name = "generate vectorized Poseidon2 trace", skip_all)] +pub fn generate_vectorized_trace_rows< + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +>( + inputs: Vec<[F; WIDTH]>, + round_constants: &RoundConstants, + extra_capacity_bits: usize, +) -> RowMajorMatrix { + let n = inputs.len(); + assert!( + n.is_multiple_of(VECTOR_LEN) && (n / VECTOR_LEN).is_power_of_two(), + "Callers expected to pad inputs to VECTOR_LEN times a power of two" + ); + + let nrows = n.div_ceil(VECTOR_LEN); + let ncols = num_cols::() + * VECTOR_LEN; + let mut vec = Vec::with_capacity((nrows * ncols) << extra_capacity_bits); + let trace = &mut vec.spare_capacity_mut()[..nrows * ncols]; + let trace = RowMajorMatrixViewMut::new(trace, ncols); + + let (prefix, perms, suffix) = unsafe { + trace.values.align_to_mut::, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >>() + }; + assert!(prefix.is_empty(), "Alignment should match"); + assert!(suffix.is_empty(), "Alignment should match"); + assert_eq!(perms.len(), n); + + perms.par_iter_mut().zip(inputs).for_each(|(perm, input)| { + generate_trace_rows_for_perm::< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(perm, input, round_constants); + }); + + unsafe { + vec.set_len(nrows * ncols); + } + + RowMajorMatrix::new(vec, ncols) +} + +// TODO: Take generic iterable +#[instrument(name = "generate Poseidon2 trace", skip_all)] +pub fn generate_trace_rows< + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>( + inputs: Vec<[F; WIDTH]>, + constants: &RoundConstants, + extra_capacity_bits: usize, +) -> RowMajorMatrix { + let n = inputs.len(); + assert!( + n.is_power_of_two(), + "Callers expected to pad inputs to a power of two" + ); + + let ncols = num_cols::(); + let mut vec = Vec::with_capacity((n * ncols) << extra_capacity_bits); + let trace = &mut vec.spare_capacity_mut()[..n * ncols]; + let trace = RowMajorMatrixViewMut::new(trace, ncols); + + let (prefix, perms, suffix) = unsafe { + trace.values.align_to_mut::, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >>() + }; + assert!(prefix.is_empty(), "Alignment should match"); + assert!(suffix.is_empty(), "Alignment should match"); + assert_eq!(perms.len(), n); + + perms.par_iter_mut().zip(inputs).for_each(|(perm, input)| { + generate_trace_rows_for_perm::< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(perm, input, constants); + }); + + unsafe { + vec.set_len(n * ncols); + } + + RowMajorMatrix::new(vec, ncols) +} + +/// `rows` will normally consist of 24 rows, with an exception for the final row. +fn generate_trace_rows_for_perm< + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>( + perm: &mut Poseidon2Cols< + MaybeUninit, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + mut state: [F; WIDTH], + constants: &RoundConstants, +) { + perm.export.write(F::ONE); + perm.inputs + .iter_mut() + .zip(state.iter()) + .for_each(|(input, &x)| { + input.write(x); + }); + + LinearLayers::external_linear_layer(&mut state); + + for (full_round, constants) in perm + .beginning_full_rounds + .iter_mut() + .zip(&constants.beginning_full_round_constants) + { + generate_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( + &mut state, full_round, constants, + ); + } + + for (partial_round, constant) in perm + .partial_rounds + .iter_mut() + .zip(&constants.partial_round_constants) + { + generate_partial_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( + &mut state, + partial_round, + *constant, + ); + } + + for (full_round, constants) in perm + .ending_full_rounds + .iter_mut() + .zip(&constants.ending_full_round_constants) + { + generate_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( + &mut state, full_round, constants, + ); + } +} + +#[inline] +fn generate_full_round< + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, +>( + state: &mut [F; WIDTH], + full_round: &mut FullRound, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>, + round_constants: &[F; WIDTH], +) { + // Combine addition of round constants and S-box application in a single loop + for ((state_i, const_i), sbox_i) in state + .iter_mut() + .zip(round_constants.iter()) + .zip(full_round.sbox.iter_mut()) + { + *state_i += *const_i; + generate_sbox(sbox_i, state_i); + } + + LinearLayers::external_linear_layer(state); + full_round + .post + .iter_mut() + .zip(*state) + .for_each(|(post, x)| { + post.write(x); + }); +} + +#[inline] +fn generate_partial_round< + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, +>( + state: &mut [F; WIDTH], + partial_round: &mut PartialRound, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>, + round_constant: F, +) { + state[0] += round_constant; + generate_sbox(&mut partial_round.sbox, &mut state[0]); + partial_round.post_sbox.write(state[0]); + LinearLayers::internal_linear_layer(state); +} + +/// Computes the S-box `x -> x^{DEGREE}` and stores the partial data required to +/// verify the computation. +/// +/// # Panics +/// +/// This method panics if the number of `REGISTERS` is not chosen optimally for the given +/// `DEGREE` or if the `DEGREE` is not supported by the S-box. The supported degrees are +/// `3`, `5`, `7`, and `11`. +#[inline] +fn generate_sbox( + sbox: &mut SBox, DEGREE, REGISTERS>, + x: &mut F, +) { + *x = match (DEGREE, REGISTERS) { + (3, 0) => x.cube(), + (5, 0) => x.exp_const_u64::<5>(), + (7, 0) => x.exp_const_u64::<7>(), + (5, 1) => { + let x2 = x.square(); + let x3 = x2 * *x; + sbox.0[0].write(x3); + x3 * x2 + } + (7, 1) => { + let x3 = x.cube(); + sbox.0[0].write(x3); + x3 * x3 * *x + } + (11, 2) => { + let x2 = x.square(); + let x3 = x2 * *x; + let x9 = x3.cube(); + sbox.0[0].write(x3); + sbox.0[1].write(x9); + x9 * x2 + } + _ => panic!( + "Unexpected (DEGREE, REGISTERS) of ({}, {})", + DEGREE, REGISTERS + ), + } +} diff --git a/poseidon2-air/src/lib.rs b/poseidon2-air/src/lib.rs new file mode 100644 index 0000000..e21683c --- /dev/null +++ b/poseidon2-air/src/lib.rs @@ -0,0 +1,17 @@ +//! And AIR for the Poseidon2 permutation. + +#![no_std] + +extern crate alloc; + +mod air; +mod columns; +mod constants; +mod generation; +mod vectorized; + +pub use air::*; +pub use columns::*; +pub use constants::*; +pub use generation::*; +pub use vectorized::*; diff --git a/poseidon2-air/src/vectorized.rs b/poseidon2-air/src/vectorized.rs new file mode 100644 index 0000000..d6fa4d0 --- /dev/null +++ b/poseidon2-air/src/vectorized.rs @@ -0,0 +1,275 @@ +use core::borrow::{Borrow, BorrowMut}; + +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_field::{PrimeCharacteristicRing, PrimeField}; +use p3_matrix::Matrix; +use p3_matrix::dense::RowMajorMatrix; +use p3_poseidon2::GenericPoseidon2LinearLayers; +use rand::distr::StandardUniform; +use rand::prelude::Distribution; +use rand::rngs::SmallRng; +use rand::{Rng, SeedableRng}; + +use crate::air::eval; +use crate::constants::RoundConstants; +use crate::{Poseidon2Air, Poseidon2Cols, generate_vectorized_trace_rows}; + +/// A "vectorized" version of Poseidon2Cols, for computing multiple Poseidon2 permutations per row. +#[repr(C)] +pub struct VectorizedPoseidon2Cols< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> { + pub(crate) cols: + [Poseidon2Cols; + VECTOR_LEN], +} + +impl< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> + Borrow< + VectorizedPoseidon2Cols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + >, + > for [T] +{ + fn borrow( + &self, + ) -> &VectorizedPoseidon2Cols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + > { + // debug_assert_eq!(self.len(), NUM_COLS); + let (prefix, shorts, suffix) = unsafe { + self.align_to::>() + }; + debug_assert!(prefix.is_empty(), "Alignment should match"); + debug_assert!(suffix.is_empty(), "Alignment should match"); + debug_assert_eq!(shorts.len(), 1); + &shorts[0] + } +} + +impl< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> + BorrowMut< + VectorizedPoseidon2Cols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + >, + > for [T] +{ + fn borrow_mut( + &mut self, + ) -> &mut VectorizedPoseidon2Cols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + > { + // debug_assert_eq!(self.len(), NUM_COLS); + let (prefix, shorts, suffix) = unsafe { + self.align_to_mut::>() + }; + debug_assert!(prefix.is_empty(), "Alignment should match"); + debug_assert!(suffix.is_empty(), "Alignment should match"); + debug_assert_eq!(shorts.len(), 1); + &mut shorts[0] + } +} + +/// A "vectorized" version of Poseidon2Air, for computing multiple Poseidon2 permutations per row. +pub struct VectorizedPoseidon2Air< + F: PrimeCharacteristicRing, + LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> { + pub(crate) air: Poseidon2Air< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, +} + +impl< + F: PrimeCharacteristicRing, + LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> + VectorizedPoseidon2Air< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + > +{ + pub const fn new( + constants: RoundConstants, + ) -> Self { + Self { + air: Poseidon2Air::new(constants), + } + } + + pub fn generate_vectorized_trace_rows( + &self, + num_hashes: usize, + extra_capacity_bits: usize, + ) -> RowMajorMatrix + where + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + StandardUniform: Distribution<[F; WIDTH]>, + { + let mut rng = SmallRng::seed_from_u64(1); + let inputs = (0..num_hashes).map(|_| rng.random()).collect(); + generate_vectorized_trace_rows::< + _, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + >(inputs, &self.air.constants, extra_capacity_bits) + } +} + +impl< + F: PrimeCharacteristicRing + Sync, + LinearLayers: Sync, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> BaseAir + for VectorizedPoseidon2Air< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + > +{ + fn width(&self) -> usize { + self.air.width() * VECTOR_LEN + } +} + +impl< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, + const VECTOR_LEN: usize, +> Air + for VectorizedPoseidon2Air< + AB::F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + > +{ + #[inline] + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0).expect("The matrix is empty?"); + let local: &VectorizedPoseidon2Cols< + _, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + VECTOR_LEN, + > = (*local).borrow(); + for perm in &local.cols { + eval(&self.air, builder, perm); + } + } +} From 9f6c0255e4ec9dec42d73bec6bc0b50f816d4880 Mon Sep 17 00:00:00 2001 From: Sai Deng Date: Wed, 8 Oct 2025 21:30:18 -0700 Subject: [PATCH 2/9] update other files --- .github/workflows/ci.yml | 1 + Cargo.toml | 3 +++ circuit-prover/Cargo.toml | 1 + 3 files changed, 5 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1da471e..7a14210 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -92,6 +92,7 @@ jobs: # cargo build --verbose --target ${{ env.target }} -p p3-fri-air # cargo build --verbose --target ${{ env.target }} -p p3-interpolation-air cargo build --verbose --target ${{ env.target }} -p p3-mmcs-air + cargo build --verbose --target ${{ env.target }} -p p3-poseidon2-air cargo build --verbose --target ${{ env.target }} -p p3-symmetric-air cargo build --verbose --target ${{ env.target }} -p p3-recursion diff --git a/Cargo.toml b/Cargo.toml index adfa802..bee81fc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "fri-air", "interpolation-air", "mmcs-air", + "poseidon2-air", "recursion", "symmetric-air", ] @@ -40,6 +41,7 @@ p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3" } p3-matrix = { git = "https://github.com/Plonky3/Plonky3" } p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3" } p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3" } +p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3" } p3-symmetric = { git = "https://github.com/Plonky3/Plonky3" } p3-uni-stark = { git = "https://github.com/Plonky3/Plonky3" } p3-util = { git = "https://github.com/Plonky3/Plonky3" } @@ -70,6 +72,7 @@ p3-field-air = { path = "field-air", version = "0.1.0" } p3-fri-air = { path = "fri-air", version = "0.1.0" } p3-interpolation-air = { path = "interpolation-air", version = "0.1.0" } p3-mmcs-air = { path = "mmcs-air", version = "0.1.0" } +p3-poseidon2-air = { path = "poseidon2-air", version = "0.1.0" } p3-recursion = { path = "recursion", version = "0.1.0" } p3-symmetric-air = { path = "symmetric-air", version = "0.1.0" } diff --git a/circuit-prover/Cargo.toml b/circuit-prover/Cargo.toml index d6f591d..1116b57 100644 --- a/circuit-prover/Cargo.toml +++ b/circuit-prover/Cargo.toml @@ -27,6 +27,7 @@ p3-matrix.workspace = true p3-maybe-rayon.workspace = true p3-merkle-tree.workspace = true p3-mmcs-air.workspace = true +p3-poseidon2-air.workspace = true p3-symmetric.workspace = true p3-uni-stark.workspace = true rand.workspace = true From f420948a080c3233ff9faae5d04e972076ddba90 Mon Sep 17 00:00:00 2001 From: hratoanina Date: Fri, 10 Oct 2025 13:56:58 +0200 Subject: [PATCH 3/9] Reuse P3 structs --- Cargo.toml | 5 +- circuit/src/tables.rs | 11 + poseidon2-air/src/air.rs | 288 ------------------ poseidon2-air/src/columns.rs | 162 ---------- poseidon2-air/src/constants.rs | 47 --- poseidon2-air/src/generation.rs | 288 ------------------ poseidon2-air/src/vectorized.rs | 275 ----------------- .../Cargo.toml | 6 +- poseidon2-circuit-air/src/air.rs | 174 +++++++++++ poseidon2-circuit-air/src/columns.rs | 133 ++++++++ .../src/lib.rs | 6 - 11 files changed, 325 insertions(+), 1070 deletions(-) delete mode 100644 poseidon2-air/src/air.rs delete mode 100644 poseidon2-air/src/columns.rs delete mode 100644 poseidon2-air/src/constants.rs delete mode 100644 poseidon2-air/src/generation.rs delete mode 100644 poseidon2-air/src/vectorized.rs rename {poseidon2-air => poseidon2-circuit-air}/Cargo.toml (84%) create mode 100644 poseidon2-circuit-air/src/air.rs create mode 100644 poseidon2-circuit-air/src/columns.rs rename {poseidon2-air => poseidon2-circuit-air}/src/lib.rs (54%) diff --git a/Cargo.toml b/Cargo.toml index bee81fc..6451259 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,7 +9,7 @@ members = [ "fri-air", "interpolation-air", "mmcs-air", - "poseidon2-air", + "poseidon2-circuit-air", "recursion", "symmetric-air", ] @@ -42,6 +42,7 @@ p3-matrix = { git = "https://github.com/Plonky3/Plonky3" } p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3" } p3-merkle-tree = { git = "https://github.com/Plonky3/Plonky3" } p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3" } +p3-poseidon2-air = { git = "https://github.com/Plonky3/Plonky3" } p3-symmetric = { git = "https://github.com/Plonky3/Plonky3" } p3-uni-stark = { git = "https://github.com/Plonky3/Plonky3" } p3-util = { git = "https://github.com/Plonky3/Plonky3" } @@ -72,7 +73,7 @@ p3-field-air = { path = "field-air", version = "0.1.0" } p3-fri-air = { path = "fri-air", version = "0.1.0" } p3-interpolation-air = { path = "interpolation-air", version = "0.1.0" } p3-mmcs-air = { path = "mmcs-air", version = "0.1.0" } -p3-poseidon2-air = { path = "poseidon2-air", version = "0.1.0" } +p3-poseidon2-circuit-air = { path = "poseidon2-circuit-air", version = "0.1.0" } p3-recursion = { path = "recursion", version = "0.1.0" } p3-symmetric-air = { path = "symmetric-air", version = "0.1.0" } diff --git a/circuit/src/tables.rs b/circuit/src/tables.rs index 539b63f..eea3399 100644 --- a/circuit/src/tables.rs +++ b/circuit/src/tables.rs @@ -120,6 +120,17 @@ pub struct MmcsPathTrace { pub final_index: Vec, } +/// Poseidon2 operation table +pub struct Poseidon2CircuitRow { + /// Inputs to the Poseidon2 permutation + pub input_values: Vec, + /// Input indices + pub input_indices: Vec, + /// Output indices + pub output_indices: Vec, +} +pub type Poseidon2CircuitTrace = Vec>; + /// Private Mmcs path data for Mmcs verification /// /// This represents the private witness information that the prover needs diff --git a/poseidon2-air/src/air.rs b/poseidon2-air/src/air.rs deleted file mode 100644 index cfa9e5b..0000000 --- a/poseidon2-air/src/air.rs +++ /dev/null @@ -1,288 +0,0 @@ -use core::borrow::Borrow; -use core::marker::PhantomData; - -use p3_air::{Air, AirBuilder, BaseAir}; -use p3_field::{PrimeCharacteristicRing, PrimeField}; -use p3_matrix::Matrix; -use p3_matrix::dense::RowMajorMatrix; -use p3_poseidon2::GenericPoseidon2LinearLayers; -use rand::distr::{Distribution, StandardUniform}; -use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; - -use crate::columns::{Poseidon2Cols, num_cols}; -use crate::constants::RoundConstants; -use crate::{FullRound, PartialRound, SBox, generate_trace_rows}; - -/// Assumes the field size is at least 16 bits. -#[derive(Debug)] -pub struct Poseidon2Air< - F: PrimeCharacteristicRing, - LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> { - pub(crate) constants: RoundConstants, - _phantom: PhantomData, -} - -impl< - F: PrimeCharacteristicRing, - LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> - Poseidon2Air< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - > -{ - pub const fn new( - constants: RoundConstants, - ) -> Self { - Self { - constants, - _phantom: PhantomData, - } - } - - pub fn generate_trace_rows( - &self, - num_hashes: usize, - extra_capacity_bits: usize, - ) -> RowMajorMatrix - where - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - StandardUniform: Distribution<[F; WIDTH]>, - { - let mut rng = SmallRng::seed_from_u64(1); - let inputs = (0..num_hashes).map(|_| rng.random()).collect(); - generate_trace_rows::< - _, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >(inputs, &self.constants, extra_capacity_bits) - } -} - -impl< - F: PrimeCharacteristicRing + Sync, - LinearLayers: Sync, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> BaseAir - for Poseidon2Air< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - > -{ - fn width(&self) -> usize { - num_cols::() - } -} - -pub(crate) fn eval< - AB: AirBuilder, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, ->( - air: &Poseidon2Air< - AB::F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >, - builder: &mut AB, - local: &Poseidon2Cols< - AB::Var, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >, -) { - let mut state: [_; WIDTH] = local.inputs.clone().map(|x| x.into()); - - LinearLayers::external_linear_layer(&mut state); - - for round in 0..HALF_FULL_ROUNDS { - eval_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( - &mut state, - &local.beginning_full_rounds[round], - &air.constants.beginning_full_round_constants[round], - builder, - ); - } - - for round in 0..PARTIAL_ROUNDS { - eval_partial_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( - &mut state, - &local.partial_rounds[round], - &air.constants.partial_round_constants[round], - builder, - ); - } - - for round in 0..HALF_FULL_ROUNDS { - eval_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( - &mut state, - &local.ending_full_rounds[round], - &air.constants.ending_full_round_constants[round], - builder, - ); - } -} - -impl< - AB: AirBuilder, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> Air - for Poseidon2Air< - AB::F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - > -{ - #[inline] - fn eval(&self, builder: &mut AB) { - let main = builder.main(); - let local = main.row_slice(0).expect("The matrix is empty?"); - let local = (*local).borrow(); - - eval::<_, _, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>( - self, builder, local, - ); - } -} - -#[inline] -fn eval_full_round< - AB: AirBuilder, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, ->( - state: &mut [AB::Expr; WIDTH], - full_round: &FullRound, - round_constants: &[AB::F; WIDTH], - builder: &mut AB, -) { - for (i, (s, r)) in state.iter_mut().zip(round_constants.iter()).enumerate() { - *s += r.clone(); - eval_sbox(&full_round.sbox[i], s, builder); - } - LinearLayers::external_linear_layer(state); - for (state_i, post_i) in state.iter_mut().zip(&full_round.post) { - builder.assert_eq(state_i.clone(), post_i.clone()); - *state_i = post_i.clone().into(); - } -} - -#[inline] -fn eval_partial_round< - AB: AirBuilder, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, ->( - state: &mut [AB::Expr; WIDTH], - partial_round: &PartialRound, - round_constant: &AB::F, - builder: &mut AB, -) { - state[0] += round_constant.clone(); - eval_sbox(&partial_round.sbox, &mut state[0], builder); - - builder.assert_eq(state[0].clone(), partial_round.post_sbox.clone()); - state[0] = partial_round.post_sbox.clone().into(); - - LinearLayers::internal_linear_layer(state); -} - -/// Evaluates the S-box over a degree-1 expression `x`. -/// -/// # Panics -/// -/// This method panics if the number of `REGISTERS` is not chosen optimally for the given -/// `DEGREE` or if the `DEGREE` is not supported by the S-box. The supported degrees are -/// `3`, `5`, `7`, and `11`. -#[inline] -fn eval_sbox( - sbox: &SBox, - x: &mut AB::Expr, - builder: &mut AB, -) where - AB: AirBuilder, -{ - *x = match (DEGREE, REGISTERS) { - (3, 0) => x.cube(), - (5, 0) => x.exp_const_u64::<5>(), - (7, 0) => x.exp_const_u64::<7>(), - (5, 1) => { - let committed_x3 = sbox.0[0].clone().into(); - let x2 = x.square(); - builder.assert_eq(committed_x3.clone(), x2.clone() * x.clone()); - committed_x3 * x2 - } - (7, 1) => { - let committed_x3 = sbox.0[0].clone().into(); - builder.assert_eq(committed_x3.clone(), x.cube()); - committed_x3.square() * x.clone() - } - (11, 2) => { - let committed_x3 = sbox.0[0].clone().into(); - let committed_x9 = sbox.0[1].clone().into(); - let x2 = x.square(); - builder.assert_eq(committed_x3.clone(), x2.clone() * x.clone()); - builder.assert_eq(committed_x9.clone(), committed_x3.cube()); - committed_x9 * x2 - } - _ => panic!( - "Unexpected (DEGREE, REGISTERS) of ({}, {})", - DEGREE, REGISTERS - ), - } -} diff --git a/poseidon2-air/src/columns.rs b/poseidon2-air/src/columns.rs deleted file mode 100644 index 082f318..0000000 --- a/poseidon2-air/src/columns.rs +++ /dev/null @@ -1,162 +0,0 @@ -use core::borrow::{Borrow, BorrowMut}; -use core::mem::size_of; - -/// Columns for a Poseidon2 AIR which computes one permutation per row. -/// -/// The columns of the STARK are divided into the three different round sections of the Poseidon2 -/// Permutation: beginning full rounds, partial rounds, and ending full rounds. For the full -/// rounds we store an [`SBox`] columnset for each state variable, and for the partial rounds we -/// store only for the first state variable. Because the matrix multiplications are linear -/// functions, we need only keep auxiliary columns for the S-box computations. -#[repr(C)] -pub struct Poseidon2Cols< - T, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> { - pub export: T, - - pub inputs: [T; WIDTH], - - /// Beginning Full Rounds - pub beginning_full_rounds: [FullRound; HALF_FULL_ROUNDS], - - /// Partial Rounds - pub partial_rounds: [PartialRound; PARTIAL_ROUNDS], - - /// Ending Full Rounds - pub ending_full_rounds: [FullRound; HALF_FULL_ROUNDS], -} - -/// Full round columns. -#[repr(C)] -pub struct FullRound { - /// Possible intermediate results within each S-box. - pub sbox: [SBox; WIDTH], - /// The post-state, i.e. the entire layer after this full round. - pub post: [T; WIDTH], -} - -/// Partial round columns. -#[repr(C)] -pub struct PartialRound -{ - /// Possible intermediate results within the S-box. - pub sbox: SBox, - /// The output of the S-box. - pub post_sbox: T, -} - -/// Possible intermediate results within an S-box. -/// -/// Use this column-set for an S-box that can be computed with `REGISTERS`-many intermediate results -/// (not counting the final output). The S-box is checked to ensure that `REGISTERS` is the optimal -/// number of registers for the given `DEGREE` for the degrees given in the Poseidon2 paper: -/// `3`, `5`, `7`, and `11`. See `eval_sbox` for more information. -#[repr(C)] -pub struct SBox(pub [T; REGISTERS]); - -pub const fn num_cols< - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, ->() -> usize { - size_of::>( - ) -} - -pub const fn make_col_map< - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, ->() -> Poseidon2Cols { - todo!() - // let indices_arr = indices_arr::< - // { num_cols::() }, - // >(); - // unsafe { - // transmute::< - // [usize; - // num_cols::()], - // Poseidon2Cols< - // usize, - // WIDTH, - // SBOX_DEGREE, - // SBOX_REGISTERS, - // HALF_FULL_ROUNDS, - // PARTIAL_ROUNDS, - // >, - // >(indices_arr) - // } -} - -impl< - T, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> Borrow> - for [T] -{ - fn borrow( - &self, - ) -> &Poseidon2Cols - { - // debug_assert_eq!(self.len(), NUM_COLS); - let (prefix, shorts, suffix) = unsafe { - self.align_to::>() - }; - debug_assert!(prefix.is_empty(), "Alignment should match"); - debug_assert!(suffix.is_empty(), "Alignment should match"); - debug_assert_eq!(shorts.len(), 1); - &shorts[0] - } -} - -impl< - T, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> BorrowMut> - for [T] -{ - fn borrow_mut( - &mut self, - ) -> &mut Poseidon2Cols - { - // debug_assert_eq!(self.len(), NUM_COLS); - let (prefix, shorts, suffix) = unsafe { - self.align_to_mut::>() - }; - debug_assert!(prefix.is_empty(), "Alignment should match"); - debug_assert!(suffix.is_empty(), "Alignment should match"); - debug_assert_eq!(shorts.len(), 1); - &mut shorts[0] - } -} diff --git a/poseidon2-air/src/constants.rs b/poseidon2-air/src/constants.rs deleted file mode 100644 index 5e50607..0000000 --- a/poseidon2-air/src/constants.rs +++ /dev/null @@ -1,47 +0,0 @@ -use p3_field::PrimeCharacteristicRing; -use rand::Rng; -use rand::distr::{Distribution, StandardUniform}; - -/// Round constants for Poseidon2, in a format that's convenient for the AIR. -#[derive(Debug, Clone)] -pub struct RoundConstants< - F: PrimeCharacteristicRing, - const WIDTH: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> { - pub(crate) beginning_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], - pub(crate) partial_round_constants: [F; PARTIAL_ROUNDS], - pub(crate) ending_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], -} - -impl< - F: PrimeCharacteristicRing, - const WIDTH: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, -> RoundConstants -{ - pub const fn new( - beginning_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], - partial_round_constants: [F; PARTIAL_ROUNDS], - ending_full_round_constants: [[F; WIDTH]; HALF_FULL_ROUNDS], - ) -> Self { - Self { - beginning_full_round_constants, - partial_round_constants, - ending_full_round_constants, - } - } - - pub fn from_rng(rng: &mut R) -> Self - where - StandardUniform: Distribution + Distribution<[F; WIDTH]>, - { - Self { - beginning_full_round_constants: core::array::from_fn(|_| rng.sample(StandardUniform)), - partial_round_constants: core::array::from_fn(|_| rng.sample(StandardUniform)), - ending_full_round_constants: core::array::from_fn(|_| rng.sample(StandardUniform)), - } - } -} diff --git a/poseidon2-air/src/generation.rs b/poseidon2-air/src/generation.rs deleted file mode 100644 index ebef843..0000000 --- a/poseidon2-air/src/generation.rs +++ /dev/null @@ -1,288 +0,0 @@ -use alloc::vec::Vec; -use core::mem::MaybeUninit; - -use p3_field::PrimeField; -use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixViewMut}; -use p3_maybe_rayon::prelude::*; -use p3_poseidon2::GenericPoseidon2LinearLayers; -use tracing::instrument; - -use crate::columns::{Poseidon2Cols, num_cols}; -use crate::{FullRound, PartialRound, RoundConstants, SBox}; - -#[instrument(name = "generate vectorized Poseidon2 trace", skip_all)] -pub fn generate_vectorized_trace_rows< - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, ->( - inputs: Vec<[F; WIDTH]>, - round_constants: &RoundConstants, - extra_capacity_bits: usize, -) -> RowMajorMatrix { - let n = inputs.len(); - assert!( - n.is_multiple_of(VECTOR_LEN) && (n / VECTOR_LEN).is_power_of_two(), - "Callers expected to pad inputs to VECTOR_LEN times a power of two" - ); - - let nrows = n.div_ceil(VECTOR_LEN); - let ncols = num_cols::() - * VECTOR_LEN; - let mut vec = Vec::with_capacity((nrows * ncols) << extra_capacity_bits); - let trace = &mut vec.spare_capacity_mut()[..nrows * ncols]; - let trace = RowMajorMatrixViewMut::new(trace, ncols); - - let (prefix, perms, suffix) = unsafe { - trace.values.align_to_mut::, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >>() - }; - assert!(prefix.is_empty(), "Alignment should match"); - assert!(suffix.is_empty(), "Alignment should match"); - assert_eq!(perms.len(), n); - - perms.par_iter_mut().zip(inputs).for_each(|(perm, input)| { - generate_trace_rows_for_perm::< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >(perm, input, round_constants); - }); - - unsafe { - vec.set_len(nrows * ncols); - } - - RowMajorMatrix::new(vec, ncols) -} - -// TODO: Take generic iterable -#[instrument(name = "generate Poseidon2 trace", skip_all)] -pub fn generate_trace_rows< - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, ->( - inputs: Vec<[F; WIDTH]>, - constants: &RoundConstants, - extra_capacity_bits: usize, -) -> RowMajorMatrix { - let n = inputs.len(); - assert!( - n.is_power_of_two(), - "Callers expected to pad inputs to a power of two" - ); - - let ncols = num_cols::(); - let mut vec = Vec::with_capacity((n * ncols) << extra_capacity_bits); - let trace = &mut vec.spare_capacity_mut()[..n * ncols]; - let trace = RowMajorMatrixViewMut::new(trace, ncols); - - let (prefix, perms, suffix) = unsafe { - trace.values.align_to_mut::, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >>() - }; - assert!(prefix.is_empty(), "Alignment should match"); - assert!(suffix.is_empty(), "Alignment should match"); - assert_eq!(perms.len(), n); - - perms.par_iter_mut().zip(inputs).for_each(|(perm, input)| { - generate_trace_rows_for_perm::< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >(perm, input, constants); - }); - - unsafe { - vec.set_len(n * ncols); - } - - RowMajorMatrix::new(vec, ncols) -} - -/// `rows` will normally consist of 24 rows, with an exception for the final row. -fn generate_trace_rows_for_perm< - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, ->( - perm: &mut Poseidon2Cols< - MaybeUninit, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >, - mut state: [F; WIDTH], - constants: &RoundConstants, -) { - perm.export.write(F::ONE); - perm.inputs - .iter_mut() - .zip(state.iter()) - .for_each(|(input, &x)| { - input.write(x); - }); - - LinearLayers::external_linear_layer(&mut state); - - for (full_round, constants) in perm - .beginning_full_rounds - .iter_mut() - .zip(&constants.beginning_full_round_constants) - { - generate_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( - &mut state, full_round, constants, - ); - } - - for (partial_round, constant) in perm - .partial_rounds - .iter_mut() - .zip(&constants.partial_round_constants) - { - generate_partial_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( - &mut state, - partial_round, - *constant, - ); - } - - for (full_round, constants) in perm - .ending_full_rounds - .iter_mut() - .zip(&constants.ending_full_round_constants) - { - generate_full_round::<_, LinearLayers, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>( - &mut state, full_round, constants, - ); - } -} - -#[inline] -fn generate_full_round< - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, ->( - state: &mut [F; WIDTH], - full_round: &mut FullRound, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>, - round_constants: &[F; WIDTH], -) { - // Combine addition of round constants and S-box application in a single loop - for ((state_i, const_i), sbox_i) in state - .iter_mut() - .zip(round_constants.iter()) - .zip(full_round.sbox.iter_mut()) - { - *state_i += *const_i; - generate_sbox(sbox_i, state_i); - } - - LinearLayers::external_linear_layer(state); - full_round - .post - .iter_mut() - .zip(*state) - .for_each(|(post, x)| { - post.write(x); - }); -} - -#[inline] -fn generate_partial_round< - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, ->( - state: &mut [F; WIDTH], - partial_round: &mut PartialRound, WIDTH, SBOX_DEGREE, SBOX_REGISTERS>, - round_constant: F, -) { - state[0] += round_constant; - generate_sbox(&mut partial_round.sbox, &mut state[0]); - partial_round.post_sbox.write(state[0]); - LinearLayers::internal_linear_layer(state); -} - -/// Computes the S-box `x -> x^{DEGREE}` and stores the partial data required to -/// verify the computation. -/// -/// # Panics -/// -/// This method panics if the number of `REGISTERS` is not chosen optimally for the given -/// `DEGREE` or if the `DEGREE` is not supported by the S-box. The supported degrees are -/// `3`, `5`, `7`, and `11`. -#[inline] -fn generate_sbox( - sbox: &mut SBox, DEGREE, REGISTERS>, - x: &mut F, -) { - *x = match (DEGREE, REGISTERS) { - (3, 0) => x.cube(), - (5, 0) => x.exp_const_u64::<5>(), - (7, 0) => x.exp_const_u64::<7>(), - (5, 1) => { - let x2 = x.square(); - let x3 = x2 * *x; - sbox.0[0].write(x3); - x3 * x2 - } - (7, 1) => { - let x3 = x.cube(); - sbox.0[0].write(x3); - x3 * x3 * *x - } - (11, 2) => { - let x2 = x.square(); - let x3 = x2 * *x; - let x9 = x3.cube(); - sbox.0[0].write(x3); - sbox.0[1].write(x9); - x9 * x2 - } - _ => panic!( - "Unexpected (DEGREE, REGISTERS) of ({}, {})", - DEGREE, REGISTERS - ), - } -} diff --git a/poseidon2-air/src/vectorized.rs b/poseidon2-air/src/vectorized.rs deleted file mode 100644 index d6fa4d0..0000000 --- a/poseidon2-air/src/vectorized.rs +++ /dev/null @@ -1,275 +0,0 @@ -use core::borrow::{Borrow, BorrowMut}; - -use p3_air::{Air, AirBuilder, BaseAir}; -use p3_field::{PrimeCharacteristicRing, PrimeField}; -use p3_matrix::Matrix; -use p3_matrix::dense::RowMajorMatrix; -use p3_poseidon2::GenericPoseidon2LinearLayers; -use rand::distr::StandardUniform; -use rand::prelude::Distribution; -use rand::rngs::SmallRng; -use rand::{Rng, SeedableRng}; - -use crate::air::eval; -use crate::constants::RoundConstants; -use crate::{Poseidon2Air, Poseidon2Cols, generate_vectorized_trace_rows}; - -/// A "vectorized" version of Poseidon2Cols, for computing multiple Poseidon2 permutations per row. -#[repr(C)] -pub struct VectorizedPoseidon2Cols< - T, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> { - pub(crate) cols: - [Poseidon2Cols; - VECTOR_LEN], -} - -impl< - T, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> - Borrow< - VectorizedPoseidon2Cols< - T, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - >, - > for [T] -{ - fn borrow( - &self, - ) -> &VectorizedPoseidon2Cols< - T, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - > { - // debug_assert_eq!(self.len(), NUM_COLS); - let (prefix, shorts, suffix) = unsafe { - self.align_to::>() - }; - debug_assert!(prefix.is_empty(), "Alignment should match"); - debug_assert!(suffix.is_empty(), "Alignment should match"); - debug_assert_eq!(shorts.len(), 1); - &shorts[0] - } -} - -impl< - T, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> - BorrowMut< - VectorizedPoseidon2Cols< - T, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - >, - > for [T] -{ - fn borrow_mut( - &mut self, - ) -> &mut VectorizedPoseidon2Cols< - T, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - > { - // debug_assert_eq!(self.len(), NUM_COLS); - let (prefix, shorts, suffix) = unsafe { - self.align_to_mut::>() - }; - debug_assert!(prefix.is_empty(), "Alignment should match"); - debug_assert!(suffix.is_empty(), "Alignment should match"); - debug_assert_eq!(shorts.len(), 1); - &mut shorts[0] - } -} - -/// A "vectorized" version of Poseidon2Air, for computing multiple Poseidon2 permutations per row. -pub struct VectorizedPoseidon2Air< - F: PrimeCharacteristicRing, - LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> { - pub(crate) air: Poseidon2Air< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >, -} - -impl< - F: PrimeCharacteristicRing, - LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> - VectorizedPoseidon2Air< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - > -{ - pub const fn new( - constants: RoundConstants, - ) -> Self { - Self { - air: Poseidon2Air::new(constants), - } - } - - pub fn generate_vectorized_trace_rows( - &self, - num_hashes: usize, - extra_capacity_bits: usize, - ) -> RowMajorMatrix - where - F: PrimeField, - LinearLayers: GenericPoseidon2LinearLayers, - StandardUniform: Distribution<[F; WIDTH]>, - { - let mut rng = SmallRng::seed_from_u64(1); - let inputs = (0..num_hashes).map(|_| rng.random()).collect(); - generate_vectorized_trace_rows::< - _, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - >(inputs, &self.air.constants, extra_capacity_bits) - } -} - -impl< - F: PrimeCharacteristicRing + Sync, - LinearLayers: Sync, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> BaseAir - for VectorizedPoseidon2Air< - F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - > -{ - fn width(&self) -> usize { - self.air.width() * VECTOR_LEN - } -} - -impl< - AB: AirBuilder, - LinearLayers: GenericPoseidon2LinearLayers, - const WIDTH: usize, - const SBOX_DEGREE: u64, - const SBOX_REGISTERS: usize, - const HALF_FULL_ROUNDS: usize, - const PARTIAL_ROUNDS: usize, - const VECTOR_LEN: usize, -> Air - for VectorizedPoseidon2Air< - AB::F, - LinearLayers, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - > -{ - #[inline] - fn eval(&self, builder: &mut AB) { - let main = builder.main(); - let local = main.row_slice(0).expect("The matrix is empty?"); - let local: &VectorizedPoseidon2Cols< - _, - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - VECTOR_LEN, - > = (*local).borrow(); - for perm in &local.cols { - eval(&self.air, builder, perm); - } - } -} diff --git a/poseidon2-air/Cargo.toml b/poseidon2-circuit-air/Cargo.toml similarity index 84% rename from poseidon2-air/Cargo.toml rename to poseidon2-circuit-air/Cargo.toml index f782c2c..ea31e5a 100644 --- a/poseidon2-air/Cargo.toml +++ b/poseidon2-circuit-air/Cargo.toml @@ -1,6 +1,6 @@ [package] -name = "p3-poseidon2-air" -description = "An AIR implementation of the Poseidon2 cryptographic hash function for use in zero-knowledge proof systems." +name = "p3-poseidon2-circuit-air" +description = "An AIR implementation of Poseidon2 modified for the circuit builder." version.workspace = true edition.workspace = true license.workspace = true @@ -11,10 +11,12 @@ categories.workspace = true [dependencies] p3-air.workspace = true +p3-circuit.workspace = true p3-field.workspace = true p3-matrix.workspace = true p3-maybe-rayon.workspace = true p3-poseidon2.workspace = true +p3-poseidon2-air.workspace = true rand.workspace = true tracing.workspace = true diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs new file mode 100644 index 0000000..52f9a9c --- /dev/null +++ b/poseidon2-circuit-air/src/air.rs @@ -0,0 +1,174 @@ +use core::borrow::Borrow; + +use p3_air::{Air, AirBuilder, BaseAir}; +use p3_circuit::tables::Poseidon2CircuitTrace; +use p3_field::{PrimeCharacteristicRing, PrimeField}; +use p3_matrix::Matrix; +use p3_matrix::dense::RowMajorMatrix; +use p3_poseidon2::GenericPoseidon2LinearLayers; +use p3_poseidon2_air::{Poseidon2Air, RoundConstants}; + +use crate::{Poseidon2CircuitCols, num_cols}; + +/// Assumes the field size is at least 16 bits. +#[derive(Debug)] +pub struct Poseidon2CircuitAir< + F: PrimeCharacteristicRing, + LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> { + p3_poseidon2: Poseidon2Air< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, +} + +impl< + F: PrimeCharacteristicRing, + LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> + Poseidon2CircuitAir< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > +{ + pub const fn new( + constants: RoundConstants, + ) -> Self { + Self { + p3_poseidon2: Poseidon2Air::new(constants), + } + } +} + +impl< + F: PrimeCharacteristicRing + Sync, + LinearLayers: Sync, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> BaseAir + for Poseidon2CircuitAir< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > +{ + fn width(&self) -> usize { + num_cols::() + } +} + +pub(crate) fn eval< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>( + air: &Poseidon2CircuitAir< + AB::F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + builder: &mut AB, + _local: &Poseidon2CircuitCols< + AB::Var, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + _next: &Poseidon2CircuitCols< + AB::Var, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, +) { + air.p3_poseidon2.eval(builder); + + // TODO: Add circuit-specific constraints. +} + +impl< + AB: AirBuilder, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> Air + for Poseidon2CircuitAir< + AB::F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > +{ + #[inline] + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.row_slice(0).expect("The matrix is empty?"); + let local = (*local).borrow(); + let next = main.row_slice(1).expect("The matrix has only one row?"); + let next = (*next).borrow(); + + eval::<_, _, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS>( + self, builder, local, next, + ); + } +} + +pub fn trace_to_matrix< + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>( + _trace: &Poseidon2CircuitTrace, +) -> &RowMajorMatrix { + // Todo: Add trace generation + todo!() +} diff --git a/poseidon2-circuit-air/src/columns.rs b/poseidon2-circuit-air/src/columns.rs new file mode 100644 index 0000000..8388d4d --- /dev/null +++ b/poseidon2-circuit-air/src/columns.rs @@ -0,0 +1,133 @@ +use core::borrow::{Borrow, BorrowMut}; + +use p3_poseidon2_air::Poseidon2Cols; + +/// Columns for a Poseidon2 AIR which computes one permutation per row. +/// +/// They extend the P3 columns with some circuit-specific columns. +#[repr(C)] +pub struct Poseidon2CircuitCols< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> { + pub poseidon2: + Poseidon2Cols, + + pub input_indices: [T; WIDTH], + pub output_indices: [T; WIDTH], +} + +pub const fn num_cols< + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +>() -> usize { + size_of::< + Poseidon2CircuitCols< + u8, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + >() +} + +impl< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> + Borrow< + Poseidon2CircuitCols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + > for [T] +{ + fn borrow( + &self, + ) -> &Poseidon2CircuitCols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > { + let (prefix, shorts, suffix) = unsafe { + self.align_to::>() + }; + debug_assert!(prefix.is_empty(), "Alignment should match"); + debug_assert!(suffix.is_empty(), "Alignment should match"); + debug_assert_eq!(shorts.len(), 1); + &shorts[0] + } +} + +impl< + T, + const WIDTH: usize, + const SBOX_DEGREE: u64, + const SBOX_REGISTERS: usize, + const HALF_FULL_ROUNDS: usize, + const PARTIAL_ROUNDS: usize, +> + BorrowMut< + Poseidon2CircuitCols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + > for [T] +{ + fn borrow_mut( + &mut self, + ) -> &mut Poseidon2CircuitCols< + T, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > { + let (prefix, shorts, suffix) = unsafe { + self.align_to_mut::>() + }; + debug_assert!(prefix.is_empty(), "Alignment should match"); + debug_assert!(suffix.is_empty(), "Alignment should match"); + debug_assert_eq!(shorts.len(), 1); + &mut shorts[0] + } +} diff --git a/poseidon2-air/src/lib.rs b/poseidon2-circuit-air/src/lib.rs similarity index 54% rename from poseidon2-air/src/lib.rs rename to poseidon2-circuit-air/src/lib.rs index e21683c..370010c 100644 --- a/poseidon2-air/src/lib.rs +++ b/poseidon2-circuit-air/src/lib.rs @@ -6,12 +6,6 @@ extern crate alloc; mod air; mod columns; -mod constants; -mod generation; -mod vectorized; pub use air::*; pub use columns::*; -pub use constants::*; -pub use generation::*; -pub use vectorized::*; From 22d9296ad89f34120f51b9826f04dfb342cacb61 Mon Sep 17 00:00:00 2001 From: hratoanina Date: Tue, 21 Oct 2025 00:08:08 +0200 Subject: [PATCH 4/9] Refactor columns --- poseidon2-circuit-air/src/air.rs | 103 ++++++++++++++++++++------- poseidon2-circuit-air/src/columns.rs | 23 +++++- 2 files changed, 99 insertions(+), 27 deletions(-) diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs index eea40ac..0fecd7a 100644 --- a/poseidon2-circuit-air/src/air.rs +++ b/poseidon2-circuit-air/src/air.rs @@ -9,11 +9,19 @@ use p3_poseidon2_air::{Poseidon2Air, RoundConstants}; use crate::{Poseidon2CircuitCols, num_cols}; +/// Extends the Poseidon2 AIR with recursion circuit-specific columns and constraints. /// Assumes the field size is at least 16 bits. +/// +/// SPECIFIC ASSUMPTIONS: +/// - Memory elements from the witness table are extension elements of degree D. +/// - RATE and CAPACITY are the number of extension elements in the rate/capacity. +/// - WIDTH is the number of field elements in the state, i.e., (RATE + CAPACITY) * D. +/// - `reset` can only be set during an absorb. #[derive(Debug)] pub struct Poseidon2CircuitAir< F: PrimeCharacteristicRing, LinearLayers, + const D: usize, const RATE: usize, const CAPACITY: usize, const WIDTH: usize, @@ -36,6 +44,7 @@ pub struct Poseidon2CircuitAir< impl< F: PrimeCharacteristicRing, LinearLayers, + const D: usize, const RATE: usize, const CAPACITY: usize, const WIDTH: usize, @@ -47,6 +56,7 @@ impl< Poseidon2CircuitAir< F, LinearLayers, + D, RATE, CAPACITY, WIDTH, @@ -59,6 +69,11 @@ impl< pub const fn new( constants: RoundConstants, ) -> Self { + assert!((CAPACITY + RATE) * D == WIDTH); + assert!(RATE.is_multiple_of(D)); + assert!(CAPACITY.is_multiple_of(D)); + assert!(WIDTH.is_multiple_of(D)); + Self { p3_poseidon2: Poseidon2Air::new(constants), } @@ -68,6 +83,7 @@ impl< impl< F: PrimeCharacteristicRing + Sync, LinearLayers: Sync, + const D: usize, const RATE: usize, const CAPACITY: usize, const WIDTH: usize, @@ -79,6 +95,7 @@ impl< for Poseidon2CircuitAir< F, LinearLayers, + D, RATE, CAPACITY, WIDTH, @@ -89,13 +106,14 @@ impl< > { fn width(&self) -> usize { - num_cols::() + num_cols::() } } pub(crate) fn eval< AB: AirBuilder, LinearLayers: GenericPoseidon2LinearLayers, + const D: usize, const RATE: usize, const CAPACITY: usize, const WIDTH: usize, @@ -107,6 +125,7 @@ pub(crate) fn eval< air: &Poseidon2CircuitAir< AB::F, LinearLayers, + D, RATE, CAPACITY, WIDTH, @@ -119,6 +138,7 @@ pub(crate) fn eval< local: &Poseidon2CircuitCols< AB::Var, WIDTH, + RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -127,6 +147,7 @@ pub(crate) fn eval< next: &Poseidon2CircuitCols< AB::Var, WIDTH, + RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -135,39 +156,71 @@ pub(crate) fn eval< ) { air.p3_poseidon2.eval(builder); - // When resetting the state, we just have to clear the capacity. The rate will be overwritten by the input. - builder - .when(local.reset.clone()) - .assert_zeros::(array::from_fn(|i| local.poseidon2.inputs[i + RATE].clone())); - - // If the next row doesn't reset, propagate the capacity. let next_no_reset = AB::Expr::ONE - next.reset.clone(); - builder - .when(next_no_reset) - .assert_zeros::(array::from_fn(|i| { - next.poseidon2.inputs[i + RATE].clone() - - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i + RATE].clone() - })); + for i in 0..(CAPACITY * D) { + // The first row has capacity zeroed. + builder + .when_first_row() + .assert_zero(local.poseidon2.inputs[RATE * D + i].clone()); + + // When resetting the state, we just have to clear the capacity. The rate will be overwritten by the input. + builder + .when(local.reset.clone()) + .assert_zero(local.poseidon2.inputs[RATE * D + i].clone()); + + // If the next row doesn't reset, propagate the capacity. + builder.when(next_no_reset.clone()).assert_zero( + next.poseidon2.inputs[RATE * D + i].clone() + - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[RATE * D + i] + .clone(), + ); + } - // If the next row is squeezing, then we propagate the entire output state forward. - let next_squeeze = AB::Expr::ONE - next.absorb.clone(); - builder - .when(next_squeeze) - .assert_zeros::(array::from_fn(|i| { - next.poseidon2.inputs[i].clone() - - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i].clone() - })); + let mut next_absorb = [AB::Expr::ZERO; RATE]; + for i in 0..RATE { + for col in next_absorb.iter_mut().take(RATE).skip(i) { + *col += next.absorb_flags[i].clone(); + } + } + let next_no_absorb = array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - next_absorb[i].clone()); + // In the next row, each rate element not being absorbed comes from the current row. + for index in 0..(RATE * D) { + let i = index / D; + let j = index % D; + builder.when(next_no_absorb[i].clone()).assert_zero( + next.poseidon2.inputs[i * D + j].clone() + - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i * D + j].clone(), + ); + } + + let mut current_absorb = [AB::Expr::ZERO; RATE]; + for i in 0..RATE { + for col in current_absorb.iter_mut().take(RATE).skip(i) { + *col += local.absorb_flags[i].clone(); + } + } + let current_no_absorb = + array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - current_absorb[i].clone()); + // During a reset, the rate elements not being absorbed are zeroed. + for (i, col) in current_no_absorb.iter().enumerate().take(RATE) { + let arr = array::from_fn::<_, D, _>(|j| local.poseidon2.inputs[i * D + j].clone().into()); + builder + .when(local.reset.clone() * col.clone()) + .assert_zeros(arr); + } + let _is_squeeze = AB::Expr::ONE - current_absorb[0].clone(); // TODO: Add all lookups: - // - If local.absorb = 1: - // * local.rate comes from input lookups. - // - If local.absorb = 0: + // - If current_absorb[i] = 1: + // * local.rate[i] comes from input lookups. + // - If is_squeeze = 1: // * local.rate is sent to output lookups. } impl< AB: AirBuilder, LinearLayers: GenericPoseidon2LinearLayers, + const D: usize, const RATE: usize, const CAPACITY: usize, const WIDTH: usize, @@ -179,6 +232,7 @@ impl< for Poseidon2CircuitAir< AB::F, LinearLayers, + D, RATE, CAPACITY, WIDTH, @@ -199,6 +253,7 @@ impl< eval::< _, _, + D, RATE, CAPACITY, WIDTH, diff --git a/poseidon2-circuit-air/src/columns.rs b/poseidon2-circuit-air/src/columns.rs index 797628e..4b39ee6 100644 --- a/poseidon2-circuit-air/src/columns.rs +++ b/poseidon2-circuit-air/src/columns.rs @@ -5,10 +5,18 @@ use p3_poseidon2_air::Poseidon2Cols; /// Columns for a Poseidon2 AIR which computes one permutation per row. /// /// They extend the P3 columns with some circuit-specific columns. +/// +/// `reset` (transparent): indicates whether the state is being reset this row. +/// `absorb_flags` (transparent): for each rate element, indicates if it is being absorbed this row. +/// At most one flag is set to 1 per row: if `absorb_flags[i]` is 1, then all elements up to the `i`-th +/// are absorbed; the rest are propagated from the previous row. +/// `input_indices` (transparent): for each rate element, indicates the index in the witness table for the +/// memory lookup. It's either received (for an absorb) or sent (for a squeeze). #[repr(C)] pub struct Poseidon2CircuitCols< T, const WIDTH: usize, + const RATE: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -18,13 +26,13 @@ pub struct Poseidon2CircuitCols< Poseidon2Cols, pub reset: T, - pub absorb: T, - pub input_indices: [T; WIDTH], - pub output_indices: [T; WIDTH], + pub absorb_flags: [T; RATE], + pub input_indices: [T; RATE], } pub const fn num_cols< const WIDTH: usize, + const RATE: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -34,6 +42,7 @@ pub const fn num_cols< Poseidon2CircuitCols< u8, WIDTH, + RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -45,6 +54,7 @@ pub const fn num_cols< impl< T, const WIDTH: usize, + const RATE: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -54,6 +64,7 @@ impl< Poseidon2CircuitCols< T, WIDTH, + RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -66,6 +77,7 @@ impl< ) -> &Poseidon2CircuitCols< T, WIDTH, + RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -75,6 +87,7 @@ impl< self.align_to:: &mut Poseidon2CircuitCols< T, WIDTH, + RATE, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -121,6 +137,7 @@ impl< self.align_to_mut:: Date: Tue, 28 Oct 2025 15:52:32 +0100 Subject: [PATCH 5/9] Add compression logic --- .github/workflows/ci.yml | 2 +- poseidon2-circuit-air/src/air.rs | 39 +++++++++++++++++++++------- poseidon2-circuit-air/src/columns.rs | 14 +++++++--- poseidon2-circuit-air/src/lib.rs | 2 +- 4 files changed, 42 insertions(+), 15 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 55ea885..636f3ad 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -97,7 +97,7 @@ jobs: # cargo build --verbose --target ${{ env.target }} -p p3-fri-air # cargo build --verbose --target ${{ env.target }} -p p3-interpolation-air cargo build --verbose --target ${{ env.target }} -p p3-mmcs-air - cargo build --verbose --target ${{ env.target }} -p p3-poseidon2-air + cargo build --verbose --target ${{ env.target }} -p p3-poseidon2-circuit-air cargo build --verbose --target ${{ env.target }} -p p3-symmetric-air cargo build --verbose --target ${{ env.target }} -p p3-recursion diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs index 0fecd7a..020236c 100644 --- a/poseidon2-circuit-air/src/air.rs +++ b/poseidon2-circuit-air/src/air.rs @@ -156,24 +156,30 @@ pub(crate) fn eval< ) { air.p3_poseidon2.eval(builder); + // SPONGE CONSTRAINTS let next_no_reset = AB::Expr::ONE - next.reset.clone(); for i in 0..(CAPACITY * D) { // The first row has capacity zeroed. builder + .when(local.is_sponge.clone()) .when_first_row() .assert_zero(local.poseidon2.inputs[RATE * D + i].clone()); // When resetting the state, we just have to clear the capacity. The rate will be overwritten by the input. builder + .when(local.is_sponge.clone()) .when(local.reset.clone()) .assert_zero(local.poseidon2.inputs[RATE * D + i].clone()); // If the next row doesn't reset, propagate the capacity. - builder.when(next_no_reset.clone()).assert_zero( - next.poseidon2.inputs[RATE * D + i].clone() - - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[RATE * D + i] - .clone(), - ); + builder + .when(local.is_sponge.clone()) + .when(next_no_reset.clone()) + .assert_zero( + next.poseidon2.inputs[RATE * D + i].clone() + - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[RATE * D + i] + .clone(), + ); } let mut next_absorb = [AB::Expr::ZERO; RATE]; @@ -187,10 +193,14 @@ pub(crate) fn eval< for index in 0..(RATE * D) { let i = index / D; let j = index % D; - builder.when(next_no_absorb[i].clone()).assert_zero( - next.poseidon2.inputs[i * D + j].clone() - - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i * D + j].clone(), - ); + builder + .when(local.is_sponge.clone()) + .when(next_no_absorb[i].clone()) + .assert_zero( + next.poseidon2.inputs[i * D + j].clone() + - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i * D + j] + .clone(), + ); } let mut current_absorb = [AB::Expr::ZERO; RATE]; @@ -201,11 +211,15 @@ pub(crate) fn eval< } let current_no_absorb = array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - current_absorb[i].clone()); + builder.assert_eq( + local.is_sponge.clone() * local.reset.clone(), + local.sponge_reset.clone(), + ); // During a reset, the rate elements not being absorbed are zeroed. for (i, col) in current_no_absorb.iter().enumerate().take(RATE) { let arr = array::from_fn::<_, D, _>(|j| local.poseidon2.inputs[i * D + j].clone().into()); builder - .when(local.reset.clone() * col.clone()) + .when(local.sponge_reset.clone() * col.clone()) .assert_zeros(arr); } @@ -215,6 +229,11 @@ pub(crate) fn eval< // * local.rate[i] comes from input lookups. // - If is_squeeze = 1: // * local.rate is sent to output lookups. + + // COMPRESSION CONSTRAINTS + // TODO: Add all lookups: + // - local input state comes from input lookups. + // - send local output state to output lookups. } impl< diff --git a/poseidon2-circuit-air/src/columns.rs b/poseidon2-circuit-air/src/columns.rs index 4b39ee6..bbe4f83 100644 --- a/poseidon2-circuit-air/src/columns.rs +++ b/poseidon2-circuit-air/src/columns.rs @@ -6,12 +6,17 @@ use p3_poseidon2_air::Poseidon2Cols; /// /// They extend the P3 columns with some circuit-specific columns. /// +/// `is_sponge` (transparent): if `1`, this row performs a sponge operation (absorb or squeeze); +/// otherwise, it performs a compression. /// `reset` (transparent): indicates whether the state is being reset this row. +/// `sponge_reset`: auxiliary column to keep constraint degrees below three. /// `absorb_flags` (transparent): for each rate element, indicates if it is being absorbed this row. /// At most one flag is set to 1 per row: if `absorb_flags[i]` is 1, then all elements up to the `i`-th /// are absorbed; the rest are propagated from the previous row. -/// `input_indices` (transparent): for each rate element, indicates the index in the witness table for the -/// memory lookup. It's either received (for an absorb) or sent (for a squeeze). +/// `input_indices` (transparent): for each input element, indicates the index in the witness table for the +/// memory lookup. It's either received (for an absorb or a compression) or sent (for a squeeze). +/// `output_indices` (transparent): for each output element, indicates the index in the witness table for the +/// memory lookup. Only used by compressions to send the output. #[repr(C)] pub struct Poseidon2CircuitCols< T, @@ -25,9 +30,12 @@ pub struct Poseidon2CircuitCols< pub poseidon2: Poseidon2Cols, + pub is_sponge: T, pub reset: T, + pub sponge_reset: T, pub absorb_flags: [T; RATE], - pub input_indices: [T; RATE], + pub input_indices: [T; WIDTH], + pub output_indices: [T; RATE], } pub const fn num_cols< diff --git a/poseidon2-circuit-air/src/lib.rs b/poseidon2-circuit-air/src/lib.rs index 370010c..f1e206b 100644 --- a/poseidon2-circuit-air/src/lib.rs +++ b/poseidon2-circuit-air/src/lib.rs @@ -1,4 +1,4 @@ -//! And AIR for the Poseidon2 permutation. +//! An AIR for the Poseidon2 table for recursion. Handles sponge operations and compressions. #![no_std] From 43c311de524eb88c1cd94a9eee9084c07bda01e6 Mon Sep 17 00:00:00 2001 From: hratoanina Date: Thu, 30 Oct 2025 11:30:20 +0100 Subject: [PATCH 6/9] temp --- poseidon2-circuit-air/src/air.rs | 115 +++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs index 020236c..c89e66a 100644 --- a/poseidon2-circuit-air/src/air.rs +++ b/poseidon2-circuit-air/src/air.rs @@ -283,3 +283,118 @@ impl< >(self, builder, local, next); } } + +#[cfg(test)] +mod test { + + use alloc::vec; + use core::array; + + use p3_baby_bear::{ + BabyBear, GenericPoseidon2LinearLayersBabyBear, default_babybear_poseidon2_16, + }; + use p3_challenger::{HashChallenger, SerializingChallenger32}; + use p3_circuit::WitnessId; + use p3_circuit::ops::MmcsVerifyConfig; + use p3_circuit::tables::{MmcsPrivateData, MmcsTrace}; + use p3_commit::ExtensionMmcs; + use p3_field::extension::BinomialExtensionField; + use p3_fri::{TwoAdicFriPcs, create_benchmark_fri_params}; + use p3_keccak::{Keccak256Hash, KeccakF}; + use p3_merkle_tree::{MerkleTreeHidingMmcs, MerkleTreeMmcs}; + use p3_poseidon2_air::RoundConstants; + use p3_symmetric::{CompressionFunctionFromHasher, PaddingFreeSponge, SerializingHasher}; + use p3_uni_stark::{StarkConfig, prove, verify}; + use rand::rngs::SmallRng; + use rand::{Rng, SeedableRng}; + + use crate::air::Poseidon2CircuitAir; + + const WIDTH: usize = 16; + const D: usize = 4; + const RATE: usize = 8; + const CAPACITY: usize = 8; + const SBOX_DEGREE: u64 = 7; + const SBOX_REGISTERS: usize = 1; + const HALF_FULL_ROUNDS: usize = 4; + const PARTIAL_ROUNDS: usize = 20; + + #[test] + fn prove_poseidon2_sponge() -> Result< + (), + p3_uni_stark::VerificationError< + p3_fri::verifier::FriError< + p3_merkle_tree::MerkleTreeError, + p3_merkle_tree::MerkleTreeError, + >, + >, + > { + type Val = BabyBear; + type Challenge = BinomialExtensionField; + + type ByteHash = Keccak256Hash; + let byte_hash = ByteHash {}; + + type U64Hash = PaddingFreeSponge; + let u64_hash = U64Hash::new(KeccakF {}); + + type FieldHash = SerializingHasher; + let field_hash = FieldHash::new(u64_hash); + + type MyCompress = CompressionFunctionFromHasher; + let compress = MyCompress::new(u64_hash); + + // WARNING: DO NOT USE SmallRng in proper applications! Use a real PRNG instead! + type ValMmcs = MerkleTreeHidingMmcs< + [Val; p3_keccak::VECTOR_LEN], + [u64; p3_keccak::VECTOR_LEN], + FieldHash, + MyCompress, + SmallRng, + 4, + 4, + >; + let mut rng = SmallRng::seed_from_u64(1); + let constants = RoundConstants::from_rng(&mut rng); + let val_mmcs = ValMmcs::new(field_hash, compress, rng); + + type ChallengeMmcs = ExtensionMmcs; + let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone()); + + type Challenger = SerializingChallenger32>; + let challenger = Challenger::from_hasher(vec![], byte_hash); + + let air: Poseidon2CircuitAir< + Val, + GenericPoseidon2LinearLayersBabyBear, + D, + RATE, + CAPACITY, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + > = Poseidon2CircuitAir::new(constants); + + // Generate random inputs. + let mut rng = SmallRng::seed_from_u64(1); + + let fri_params = create_benchmark_fri_params(challenge_mmcs); + + let trace = air.generate_trace_rows(NUM_PERMUTATIONS, fri_params.log_blowup); + + type Dft = p3_dft::Radix2Bowers; + let dft = Dft::default(); + + type Pcs = TwoAdicFriPcs; + let pcs = Pcs::new(dft, val_mmcs, fri_params); + + type MyConfig = StarkConfig; + let config = MyConfig::new(pcs, challenger); + + let proof = prove(&config, &air, trace, &vec![]); + + verify(&config, &air, &proof, &vec![]) + } +} From 71aafa55cf4b81ce174d515b4511e657c62749e2 Mon Sep 17 00:00:00 2001 From: hratoanina Date: Tue, 4 Nov 2025 00:25:37 +0100 Subject: [PATCH 7/9] alignment wip --- circuit/src/tables.rs | 3 + circuit/src/tables/poseidon2.rs | 8 + poseidon2-circuit-air/Cargo.toml | 3 +- poseidon2-circuit-air/src/air.rs | 217 +++++++++++++++++++++++++-- poseidon2-circuit-air/src/columns.rs | 1 + 5 files changed, 219 insertions(+), 13 deletions(-) diff --git a/circuit/src/tables.rs b/circuit/src/tables.rs index ff522e1..024a7e0 100644 --- a/circuit/src/tables.rs +++ b/circuit/src/tables.rs @@ -11,6 +11,9 @@ use crate::{CircuitError, CircuitField}; mod mmcs; pub use mmcs::{MmcsPathTrace, MmcsPrivateData, MmcsTrace}; +mod poseidon2; +pub use poseidon2::{Poseidon2CircuitRow, Poseidon2CircuitTrace}; + /// Execution traces for all tables #[derive(Debug, Clone)] pub struct Traces { diff --git a/circuit/src/tables/poseidon2.rs b/circuit/src/tables/poseidon2.rs index f0b90ff..8c47e3b 100644 --- a/circuit/src/tables/poseidon2.rs +++ b/circuit/src/tables/poseidon2.rs @@ -1,5 +1,13 @@ +use alloc::vec::Vec; + /// Poseidon2 operation table pub struct Poseidon2CircuitRow { + /// Poseidon2 operation type + pub is_sponge: bool, + /// Reset flag + pub reset: bool, + /// Absorb flags + pub absorb_flags: Vec, /// Inputs to the Poseidon2 permutation pub input_values: Vec, /// Input indices diff --git a/poseidon2-circuit-air/Cargo.toml b/poseidon2-circuit-air/Cargo.toml index ea31e5a..1a01944 100644 --- a/poseidon2-circuit-air/Cargo.toml +++ b/poseidon2-circuit-air/Cargo.toml @@ -17,7 +17,9 @@ p3-matrix.workspace = true p3-maybe-rayon.workspace = true p3-poseidon2.workspace = true p3-poseidon2-air.workspace = true +p3-symmetric.workspace = true +itertools.workspace = true rand.workspace = true tracing.workspace = true @@ -33,7 +35,6 @@ p3-fri.workspace = true p3-keccak.workspace = true p3-koala-bear.workspace = true p3-merkle-tree.workspace = true -p3-symmetric.workspace = true p3-uni-stark.workspace = true tracing-forest = { workspace = true, features = ["ansi", "smallvec"] } diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs index c89e66a..437508c 100644 --- a/poseidon2-circuit-air/src/air.rs +++ b/poseidon2-circuit-air/src/air.rs @@ -1,11 +1,20 @@ -use core::array; +use alloc::string::String; +use alloc::vec; +use alloc::vec::Vec; use core::borrow::Borrow; +use core::mem::MaybeUninit; +use core::{array, num}; +use itertools::izip; use p3_air::{Air, AirBuilder, BaseAir}; -use p3_field::PrimeCharacteristicRing; +use p3_circuit::tables::{Poseidon2CircuitRow, Poseidon2CircuitTrace}; +use p3_field::{PrimeCharacteristicRing, PrimeField}; use p3_matrix::Matrix; +use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixViewMut}; use p3_poseidon2::GenericPoseidon2LinearLayers; -use p3_poseidon2_air::{Poseidon2Air, RoundConstants}; +use p3_poseidon2_air::{Poseidon2Air, RoundConstants, generate_trace_rows}; +use p3_symmetric::CryptographicPermutation; +use tracing::info; use crate::{Poseidon2CircuitCols, num_cols}; @@ -42,8 +51,8 @@ pub struct Poseidon2CircuitAir< } impl< - F: PrimeCharacteristicRing, - LinearLayers, + F: PrimeField, + LinearLayers: GenericPoseidon2LinearLayers, const D: usize, const RATE: usize, const CAPACITY: usize, @@ -70,14 +79,152 @@ impl< constants: RoundConstants, ) -> Self { assert!((CAPACITY + RATE) * D == WIDTH); - assert!(RATE.is_multiple_of(D)); - assert!(CAPACITY.is_multiple_of(D)); assert!(WIDTH.is_multiple_of(D)); Self { p3_poseidon2: Poseidon2Air::new(constants), } } + + pub fn generate_trace_rows>( + &self, + sponge_ops: Poseidon2CircuitTrace, + constants: &RoundConstants, + extra_capacity_bits: usize, + perm: P, + ) -> RowMajorMatrix { + let n = sponge_ops.len(); + assert!( + n.is_power_of_two(), + "Callers expected to pad inputs to a power of two" + ); + + let num_circuit_cols = 3 + 2 * RATE + WIDTH; + // let mut circuit_trace = Vec::with_capacity(n * num_circuit_cols); + let mut circuit_trace = vec![F::ZERO; n * num_circuit_cols]; + let mut circuit_trace = RowMajorMatrixViewMut::new(&mut circuit_trace, num_circuit_cols); + + let mut state = [F::ZERO; WIDTH]; + let mut inputs = Vec::with_capacity(n); + for (i, op) in sponge_ops.iter().enumerate() { + info!("i: {i}"); + let Poseidon2CircuitRow { + is_sponge, + reset, + absorb_flags, + input_values, + input_indices, + output_indices, + } = op; + info!("hmm"); + + let row = circuit_trace.row_mut(i); + + row[0] = if *is_sponge { F::ONE } else { F::ZERO }; + row[1] = if *reset { F::ONE } else { F::ZERO }; + row[2] = if *is_sponge && *reset { + F::ONE + } else { + F::ZERO + }; + for j in 0..RATE { + row[3 + j] = if absorb_flags[j] { F::ONE } else { F::ZERO }; + } + for j in 0..RATE { + row[3 + RATE + j] = F::from_u32(input_indices[j]); + } + for j in 0..RATE { + row[3 + RATE + WIDTH + j] = F::from_u32(output_indices[j]); + } + + let mut index_absorb = [false; RATE]; + for j in 0..RATE { + if absorb_flags[j] { + for k in 0..=j { + index_absorb[k] = true; + } + } + } + + for j in 0..RATE { + if index_absorb[j] { + for d in 0..D { + let idx = j * D + d; + state[idx] = input_values[idx]; + } + } + } + + if *reset || !*is_sponge { + // Compression or reset: reset capacity + for j in 0..(CAPACITY * D) { + state[RATE * D + j] = F::ZERO; + } + + inputs.push(state.clone()); + state = perm.permute(state); + } + } + + let p2_trace = generate_trace_rows::< + F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(inputs, constants, extra_capacity_bits); + + let ncols = + num_cols::( + ); + + info!("ncols: {}", ncols); + + let p2_ncols = p3_poseidon2_air::num_cols::< + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(); + info!("p2_ncols: {}", p2_ncols); + + let mut vec = Vec::with_capacity((n * ncols) << extra_capacity_bits); + let trace = &mut vec.spare_capacity_mut()[..n * ncols]; + let trace = RowMajorMatrixViewMut::new(trace, ncols); + + let (prefix, perms, suffix) = unsafe { + trace.values.align_to_mut::, + WIDTH, + RATE, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >>() + }; + assert!(prefix.is_empty(), "Alignment should match"); + assert!(suffix.is_empty(), "Alignment should match"); + assert_eq!(perms.len(), n); + + for (row, circuit_row, perm_row) in izip!(perms, circuit_trace.rows(), p2_trace.rows()) { + let left_part = circuit_row.collect::>(); + let right_part = perm_row.collect::>(); + // row[..num_circuit_cols].copy_from_slice(&left_part); + // row[num_circuit_cols..].copy_from_slice(&right_part); + } + + info!("partial trace: {:?}", vec); + + unsafe { + vec.set_len(n * ncols); + } + + RowMajorMatrix::new(vec, ncols) + } } impl< @@ -296,7 +443,9 @@ mod test { use p3_challenger::{HashChallenger, SerializingChallenger32}; use p3_circuit::WitnessId; use p3_circuit::ops::MmcsVerifyConfig; - use p3_circuit::tables::{MmcsPrivateData, MmcsTrace}; + use p3_circuit::tables::{ + MmcsPrivateData, MmcsTrace, Poseidon2CircuitRow, Poseidon2CircuitTrace, + }; use p3_commit::ExtensionMmcs; use p3_field::extension::BinomialExtensionField; use p3_fri::{TwoAdicFriPcs, create_benchmark_fri_params}; @@ -307,18 +456,42 @@ mod test { use p3_uni_stark::{StarkConfig, prove, verify}; use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; + use tracing_forest::ForestLayer; + use tracing_forest::util::LevelFilter; + use tracing_subscriber::layer::SubscriberExt; + use tracing_subscriber::util::SubscriberInitExt; + use tracing_subscriber::{EnvFilter, Registry}; use crate::air::Poseidon2CircuitAir; const WIDTH: usize = 16; const D: usize = 4; - const RATE: usize = 8; - const CAPACITY: usize = 8; + const RATE: usize = 2; + const CAPACITY: usize = 2; const SBOX_DEGREE: u64 = 7; const SBOX_REGISTERS: usize = 1; const HALF_FULL_ROUNDS: usize = 4; const PARTIAL_ROUNDS: usize = 20; + const P2_NUM_COLS: usize = p3_poseidon2_air::num_cols::< + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(); + + fn init_logger() { + let env_filter = EnvFilter::builder() + .with_default_directive(LevelFilter::INFO.into()) + .from_env_lossy(); + + Registry::default() + .with(env_filter) + .with(ForestLayer::default()) + .init(); + } + #[test] fn prove_poseidon2_sponge() -> Result< (), @@ -329,6 +502,7 @@ mod test { >, >, > { + init_logger(); type Val = BabyBear; type Challenge = BinomialExtensionField; @@ -375,14 +549,33 @@ mod test { SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS, - > = Poseidon2CircuitAir::new(constants); + > = Poseidon2CircuitAir::new(constants.clone()); // Generate random inputs. let mut rng = SmallRng::seed_from_u64(1); + let a: Poseidon2CircuitRow = Poseidon2CircuitRow { + is_sponge: true, + reset: true, + absorb_flags: vec![false, true], + input_values: (0..RATE * D).map(|_| rng.random()).collect(), + input_indices: vec![0; RATE], + output_indices: vec![0; RATE], + }; + + let b: Poseidon2CircuitRow = Poseidon2CircuitRow { + is_sponge: true, + reset: true, + absorb_flags: vec![false, true], + input_values: (0..RATE * D).map(|_| rng.random()).collect(), + input_indices: vec![0; RATE], + output_indices: vec![0; RATE], + }; + let fri_params = create_benchmark_fri_params(challenge_mmcs); - let trace = air.generate_trace_rows(NUM_PERMUTATIONS, fri_params.log_blowup); + let perm = default_babybear_poseidon2_16(); + let trace = air.generate_trace_rows(vec![a, b], &constants, fri_params.log_blowup, perm); type Dft = p3_dft::Radix2Bowers; let dft = Dft::default(); diff --git a/poseidon2-circuit-air/src/columns.rs b/poseidon2-circuit-air/src/columns.rs index bbe4f83..606ccb0 100644 --- a/poseidon2-circuit-air/src/columns.rs +++ b/poseidon2-circuit-air/src/columns.rs @@ -36,6 +36,7 @@ pub struct Poseidon2CircuitCols< pub absorb_flags: [T; RATE], pub input_indices: [T; WIDTH], pub output_indices: [T; RATE], + pub _padding: [T; 290], } pub const fn num_cols< From d842bbdc050ecfee58d1b0c28f72faa56f5dca27 Mon Sep 17 00:00:00 2001 From: hratoanina Date: Thu, 6 Nov 2025 02:05:00 +0100 Subject: [PATCH 8/9] Add sub builder --- poseidon2-circuit-air/src/air.rs | 357 ++++++++++++++--------- poseidon2-circuit-air/src/columns.rs | 40 ++- poseidon2-circuit-air/src/lib.rs | 1 + poseidon2-circuit-air/src/sub_builder.rs | 108 +++++++ 4 files changed, 355 insertions(+), 151 deletions(-) create mode 100644 poseidon2-circuit-air/src/sub_builder.rs diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs index 437508c..2bb82d0 100644 --- a/poseidon2-circuit-air/src/air.rs +++ b/poseidon2-circuit-air/src/air.rs @@ -1,11 +1,8 @@ -use alloc::string::String; use alloc::vec; use alloc::vec::Vec; +use core::array; use core::borrow::Borrow; -use core::mem::MaybeUninit; -use core::{array, num}; -use itertools::izip; use p3_air::{Air, AirBuilder, BaseAir}; use p3_circuit::tables::{Poseidon2CircuitRow, Poseidon2CircuitTrace}; use p3_field::{PrimeCharacteristicRing, PrimeField}; @@ -14,8 +11,8 @@ use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixViewMut}; use p3_poseidon2::GenericPoseidon2LinearLayers; use p3_poseidon2_air::{Poseidon2Air, RoundConstants, generate_trace_rows}; use p3_symmetric::CryptographicPermutation; -use tracing::info; +use crate::sub_builder::SubAirBuilder; use crate::{Poseidon2CircuitCols, num_cols}; /// Extends the Poseidon2 AIR with recursion circuit-specific columns and constraints. @@ -31,9 +28,10 @@ pub struct Poseidon2CircuitAir< F: PrimeCharacteristicRing, LinearLayers, const D: usize, - const RATE: usize, - const CAPACITY: usize, const WIDTH: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, + const CAPACITY_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -54,9 +52,10 @@ impl< F: PrimeField, LinearLayers: GenericPoseidon2LinearLayers, const D: usize, - const RATE: usize, - const CAPACITY: usize, const WIDTH: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, + const CAPACITY_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -66,9 +65,10 @@ impl< F, LinearLayers, D, - RATE, - CAPACITY, WIDTH, + WIDTH_EXT, + RATE_EXT, + CAPACITY_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -78,8 +78,8 @@ impl< pub const fn new( constants: RoundConstants, ) -> Self { - assert!((CAPACITY + RATE) * D == WIDTH); - assert!(WIDTH.is_multiple_of(D)); + assert!(CAPACITY_EXT + RATE_EXT == WIDTH_EXT); + assert!(WIDTH_EXT * D == WIDTH); Self { p3_poseidon2: Poseidon2Air::new(constants), @@ -99,15 +99,13 @@ impl< "Callers expected to pad inputs to a power of two" ); - let num_circuit_cols = 3 + 2 * RATE + WIDTH; - // let mut circuit_trace = Vec::with_capacity(n * num_circuit_cols); + let num_circuit_cols = 3 + 2 * RATE_EXT + WIDTH_EXT; let mut circuit_trace = vec![F::ZERO; n * num_circuit_cols]; let mut circuit_trace = RowMajorMatrixViewMut::new(&mut circuit_trace, num_circuit_cols); let mut state = [F::ZERO; WIDTH]; let mut inputs = Vec::with_capacity(n); for (i, op) in sponge_ops.iter().enumerate() { - info!("i: {i}"); let Poseidon2CircuitRow { is_sponge, reset, @@ -116,7 +114,6 @@ impl< input_indices, output_indices, } = op; - info!("hmm"); let row = circuit_trace.row_mut(i); @@ -127,18 +124,18 @@ impl< } else { F::ZERO }; - for j in 0..RATE { + for j in 0..RATE_EXT { row[3 + j] = if absorb_flags[j] { F::ONE } else { F::ZERO }; } - for j in 0..RATE { - row[3 + RATE + j] = F::from_u32(input_indices[j]); + for j in 0..RATE_EXT { + row[3 + RATE_EXT + j] = F::from_u32(input_indices[j]); } - for j in 0..RATE { - row[3 + RATE + WIDTH + j] = F::from_u32(output_indices[j]); + for j in 0..RATE_EXT { + row[3 + RATE_EXT + WIDTH_EXT + j] = F::from_u32(output_indices[j]); } - let mut index_absorb = [false; RATE]; - for j in 0..RATE { + let mut index_absorb = [false; RATE_EXT]; + for j in 0..RATE_EXT { if absorb_flags[j] { for k in 0..=j { index_absorb[k] = true; @@ -146,24 +143,32 @@ impl< } } - for j in 0..RATE { + for j in 0..RATE_EXT { if index_absorb[j] { for d in 0..D { let idx = j * D + d; state[idx] = input_values[idx]; } + } else { + if *reset { + // During a reset, non-absorbed rate elements are zeroed. + for d in 0..D { + let idx = j * D + d; + state[idx] = F::ZERO; + } + } } } if *reset || !*is_sponge { // Compression or reset: reset capacity - for j in 0..(CAPACITY * D) { - state[RATE * D + j] = F::ZERO; + for j in 0..(CAPACITY_EXT * D) { + state[RATE_EXT * D + j] = F::ZERO; } - - inputs.push(state.clone()); - state = perm.permute(state); } + + inputs.push(state.clone()); + state = perm.permute(state); } let p2_trace = generate_trace_rows::< @@ -176,11 +181,7 @@ impl< PARTIAL_ROUNDS, >(inputs, constants, extra_capacity_bits); - let ncols = - num_cols::( - ); - - info!("ncols: {}", ncols); + let ncols = self.width(); let p2_ncols = p3_poseidon2_air::num_cols::< WIDTH, @@ -189,35 +190,24 @@ impl< HALF_FULL_ROUNDS, PARTIAL_ROUNDS, >(); - info!("p2_ncols: {}", p2_ncols); - - let mut vec = Vec::with_capacity((n * ncols) << extra_capacity_bits); - let trace = &mut vec.spare_capacity_mut()[..n * ncols]; - let trace = RowMajorMatrixViewMut::new(trace, ncols); - - let (prefix, perms, suffix) = unsafe { - trace.values.align_to_mut::, - WIDTH, - RATE, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >>() - }; - assert!(prefix.is_empty(), "Alignment should match"); - assert!(suffix.is_empty(), "Alignment should match"); - assert_eq!(perms.len(), n); - - for (row, circuit_row, perm_row) in izip!(perms, circuit_trace.rows(), p2_trace.rows()) { - let left_part = circuit_row.collect::>(); - let right_part = perm_row.collect::>(); - // row[..num_circuit_cols].copy_from_slice(&left_part); - // row[num_circuit_cols..].copy_from_slice(&right_part); - } - info!("partial trace: {:?}", vec); + let mut vec = vec![F::ZERO; n * ncols]; + + for i in 0..n { + let row = &mut vec[(i * ncols)..((i + 1) * ncols)]; + let left_part = p2_trace + .row(i) + .expect("Missing row {i}?") + .into_iter() + .collect::>(); + let right_part = circuit_trace + .row(i) + .expect("Missing row {i}?") + .into_iter() + .collect::>(); + row[..p2_ncols].copy_from_slice(&left_part); + row[p2_ncols..].copy_from_slice(&right_part); + } unsafe { vec.set_len(n * ncols); @@ -231,9 +221,10 @@ impl< F: PrimeCharacteristicRing + Sync, LinearLayers: Sync, const D: usize, - const RATE: usize, - const CAPACITY: usize, const WIDTH: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, + const CAPACITY_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -243,9 +234,10 @@ impl< F, LinearLayers, D, - RATE, - CAPACITY, WIDTH, + WIDTH_EXT, + RATE_EXT, + CAPACITY_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -253,7 +245,15 @@ impl< > { fn width(&self) -> usize { - num_cols::() + num_cols::< + WIDTH, + WIDTH_EXT, + RATE_EXT, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >() } } @@ -261,9 +261,10 @@ pub(crate) fn eval< AB: AirBuilder, LinearLayers: GenericPoseidon2LinearLayers, const D: usize, - const RATE: usize, - const CAPACITY: usize, const WIDTH: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, + const CAPACITY_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -273,9 +274,10 @@ pub(crate) fn eval< AB::F, LinearLayers, D, - RATE, - CAPACITY, WIDTH, + WIDTH_EXT, + RATE_EXT, + CAPACITY_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -285,7 +287,8 @@ pub(crate) fn eval< local: &Poseidon2CircuitCols< AB::Var, WIDTH, - RATE, + WIDTH_EXT, + RATE_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -294,55 +297,61 @@ pub(crate) fn eval< next: &Poseidon2CircuitCols< AB::Var, WIDTH, - RATE, + WIDTH_EXT, + RATE_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS, >, ) { - air.p3_poseidon2.eval(builder); - // SPONGE CONSTRAINTS let next_no_reset = AB::Expr::ONE - next.reset.clone(); - for i in 0..(CAPACITY * D) { + for i in 0..(CAPACITY_EXT * D) { // The first row has capacity zeroed. builder .when(local.is_sponge.clone()) .when_first_row() - .assert_zero(local.poseidon2.inputs[RATE * D + i].clone()); + .assert_zero(local.poseidon2.inputs[RATE_EXT * D + i].clone()); // When resetting the state, we just have to clear the capacity. The rate will be overwritten by the input. builder .when(local.is_sponge.clone()) .when(local.reset.clone()) - .assert_zero(local.poseidon2.inputs[RATE * D + i].clone()); + .assert_zero(local.poseidon2.inputs[RATE_EXT * D + i].clone()); // If the next row doesn't reset, propagate the capacity. builder - .when(local.is_sponge.clone()) + .when_transition() + .when(next.is_sponge.clone()) .when(next_no_reset.clone()) .assert_zero( - next.poseidon2.inputs[RATE * D + i].clone() - - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[RATE * D + i] + next.poseidon2.inputs[RATE_EXT * D + i].clone() + - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post + [RATE_EXT * D + i] .clone(), ); } - let mut next_absorb = [AB::Expr::ZERO; RATE]; - for i in 0..RATE { - for col in next_absorb.iter_mut().take(RATE).skip(i) { + let mut next_absorb = [AB::Expr::ZERO; RATE_EXT]; + for i in 0..RATE_EXT { + for col in next_absorb.iter_mut().take(i + 1) { *col += next.absorb_flags[i].clone(); } } - let next_no_absorb = array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - next_absorb[i].clone()); - // In the next row, each rate element not being absorbed comes from the current row. - for index in 0..(RATE * D) { + let next_no_absorb = + array::from_fn::<_, RATE_EXT, _>(|i| AB::Expr::ONE - next_absorb[i].clone()); + // In the next row, each rate element not being absorbed is either: + // - zeroed if the next row is a reset (handled elsewhere); + // - copied from the current row if the next row is not a reset. + // We omit the `is_sponge` check because in a compression all absorb flags are set. + for index in 0..(RATE_EXT * D) { let i = index / D; let j = index % D; builder - .when(local.is_sponge.clone()) + .when_transition() .when(next_no_absorb[i].clone()) + .when(next_no_reset.clone()) .assert_zero( next.poseidon2.inputs[i * D + j].clone() - local.poseidon2.ending_full_rounds[HALF_FULL_ROUNDS - 1].post[i * D + j] @@ -350,27 +359,27 @@ pub(crate) fn eval< ); } - let mut current_absorb = [AB::Expr::ZERO; RATE]; - for i in 0..RATE { - for col in current_absorb.iter_mut().take(RATE).skip(i) { + let mut current_absorb = [AB::Expr::ZERO; RATE_EXT]; + for i in 0..RATE_EXT { + for col in current_absorb.iter_mut().take(i + 1) { *col += local.absorb_flags[i].clone(); } } let current_no_absorb = - array::from_fn::<_, RATE, _>(|i| AB::Expr::ONE - current_absorb[i].clone()); + array::from_fn::<_, RATE_EXT, _>(|i| AB::Expr::ONE - current_absorb[i].clone()); builder.assert_eq( local.is_sponge.clone() * local.reset.clone(), local.sponge_reset.clone(), ); // During a reset, the rate elements not being absorbed are zeroed. - for (i, col) in current_no_absorb.iter().enumerate().take(RATE) { + for (i, col) in current_no_absorb.iter().enumerate() { let arr = array::from_fn::<_, D, _>(|j| local.poseidon2.inputs[i * D + j].clone().into()); builder .when(local.sponge_reset.clone() * col.clone()) .assert_zeros(arr); } - let _is_squeeze = AB::Expr::ONE - current_absorb[0].clone(); + // let _is_squeeze = AB::Expr::ONE - current_absorb[0].clone(); // TODO: Add all lookups: // - If current_absorb[i] = 1: // * local.rate[i] comes from input lookups. @@ -381,15 +390,40 @@ pub(crate) fn eval< // TODO: Add all lookups: // - local input state comes from input lookups. // - send local output state to output lookups. + + let p3_poseidon2_num_cols = p3_poseidon2_air::num_cols::< + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >(); + let mut sub_builder = SubAirBuilder::< + AB, + Poseidon2Air< + AB::F, + LinearLayers, + WIDTH, + SBOX_DEGREE, + SBOX_REGISTERS, + HALF_FULL_ROUNDS, + PARTIAL_ROUNDS, + >, + AB::Var, + >::new(builder, 0..p3_poseidon2_num_cols); + + // Eval the Plonky3 Poseidon2 air. + air.p3_poseidon2.eval(&mut sub_builder); } impl< AB: AirBuilder, LinearLayers: GenericPoseidon2LinearLayers, const D: usize, - const RATE: usize, - const CAPACITY: usize, const WIDTH: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, + const CAPACITY_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -399,9 +433,10 @@ impl< AB::F, LinearLayers, D, - RATE, - CAPACITY, WIDTH, + WIDTH_EXT, + RATE_EXT, + CAPACITY_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -420,9 +455,10 @@ impl< _, _, D, - RATE, - CAPACITY, WIDTH, + WIDTH_EXT, + RATE_EXT, + CAPACITY_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -435,22 +471,19 @@ impl< mod test { use alloc::vec; - use core::array; use p3_baby_bear::{ - BabyBear, GenericPoseidon2LinearLayersBabyBear, default_babybear_poseidon2_16, + BabyBear, GenericPoseidon2LinearLayersBabyBear, Poseidon2ExternalLayerBabyBear, + Poseidon2InternalLayerBabyBear, }; use p3_challenger::{HashChallenger, SerializingChallenger32}; - use p3_circuit::WitnessId; - use p3_circuit::ops::MmcsVerifyConfig; - use p3_circuit::tables::{ - MmcsPrivateData, MmcsTrace, Poseidon2CircuitRow, Poseidon2CircuitTrace, - }; + use p3_circuit::tables::Poseidon2CircuitRow; use p3_commit::ExtensionMmcs; use p3_field::extension::BinomialExtensionField; use p3_fri::{TwoAdicFriPcs, create_benchmark_fri_params}; use p3_keccak::{Keccak256Hash, KeccakF}; - use p3_merkle_tree::{MerkleTreeHidingMmcs, MerkleTreeMmcs}; + use p3_merkle_tree::MerkleTreeHidingMmcs; + use p3_poseidon2::{ExternalLayerConstants, Poseidon2}; use p3_poseidon2_air::RoundConstants; use p3_symmetric::{CompressionFunctionFromHasher, PaddingFreeSponge, SerializingHasher}; use p3_uni_stark::{StarkConfig, prove, verify}; @@ -464,23 +497,16 @@ mod test { use crate::air::Poseidon2CircuitAir; - const WIDTH: usize = 16; const D: usize = 4; - const RATE: usize = 2; - const CAPACITY: usize = 2; + const WIDTH: usize = 16; + const WIDTH_EXT: usize = 4; + const RATE_EXT: usize = 2; + const CAPACITY_EXT: usize = 2; const SBOX_DEGREE: u64 = 7; const SBOX_REGISTERS: usize = 1; const HALF_FULL_ROUNDS: usize = 4; const PARTIAL_ROUNDS: usize = 20; - const P2_NUM_COLS: usize = p3_poseidon2_air::num_cols::< - WIDTH, - SBOX_DEGREE, - SBOX_REGISTERS, - HALF_FULL_ROUNDS, - PARTIAL_ROUNDS, - >(); - fn init_logger() { let env_filter = EnvFilter::builder() .with_default_directive(LevelFilter::INFO.into()) @@ -529,8 +555,7 @@ mod test { 4, >; let mut rng = SmallRng::seed_from_u64(1); - let constants = RoundConstants::from_rng(&mut rng); - let val_mmcs = ValMmcs::new(field_hash, compress, rng); + let val_mmcs = ValMmcs::new(field_hash, compress, rng.clone()); type ChallengeMmcs = ExtensionMmcs; let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone()); @@ -538,13 +563,40 @@ mod test { type Challenger = SerializingChallenger32>; let challenger = Challenger::from_hasher(vec![], byte_hash); + let fri_params = create_benchmark_fri_params(challenge_mmcs); + + let beginning_full_constants = rng.random(); + let partial_constants = rng.random(); + let ending_full_constants = rng.random(); + + let constants = RoundConstants::new( + beginning_full_constants, + partial_constants, + ending_full_constants, + ); + + let perm = Poseidon2::< + Val, + Poseidon2ExternalLayerBabyBear, + Poseidon2InternalLayerBabyBear, + WIDTH, + SBOX_DEGREE, + >::new( + ExternalLayerConstants::new( + beginning_full_constants.to_vec(), + ending_full_constants.to_vec(), + ), + partial_constants.to_vec(), + ); + let air: Poseidon2CircuitAir< Val, GenericPoseidon2LinearLayersBabyBear, D, - RATE, - CAPACITY, WIDTH, + WIDTH_EXT, + RATE_EXT, + CAPACITY_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -554,28 +606,61 @@ mod test { // Generate random inputs. let mut rng = SmallRng::seed_from_u64(1); - let a: Poseidon2CircuitRow = Poseidon2CircuitRow { + // Absorb + let sponge_a: Poseidon2CircuitRow = Poseidon2CircuitRow { is_sponge: true, reset: true, absorb_flags: vec![false, true], - input_values: (0..RATE * D).map(|_| rng.random()).collect(), - input_indices: vec![0; RATE], - output_indices: vec![0; RATE], + input_values: (0..RATE_EXT * D).map(|_| rng.random()).collect(), + input_indices: vec![0; RATE_EXT], + output_indices: vec![0; RATE_EXT], }; - let b: Poseidon2CircuitRow = Poseidon2CircuitRow { + // Absorb + let sponge_b: Poseidon2CircuitRow = Poseidon2CircuitRow { is_sponge: true, - reset: true, + reset: false, absorb_flags: vec![false, true], - input_values: (0..RATE * D).map(|_| rng.random()).collect(), - input_indices: vec![0; RATE], - output_indices: vec![0; RATE], + input_values: (0..RATE_EXT * D).map(|_| rng.random()).collect(), + input_indices: vec![0; RATE_EXT], + output_indices: vec![0; RATE_EXT], }; - let fri_params = create_benchmark_fri_params(challenge_mmcs); + // Squeeze + let sponge_c: Poseidon2CircuitRow = Poseidon2CircuitRow { + is_sponge: true, + reset: false, + absorb_flags: vec![false, false], + input_values: vec![Val::new(0); RATE_EXT * D], + input_indices: vec![0; RATE_EXT], + output_indices: vec![0; RATE_EXT], + }; - let perm = default_babybear_poseidon2_16(); - let trace = air.generate_trace_rows(vec![a, b], &constants, fri_params.log_blowup, perm); + // Absorb one element with reset + let sponge_d: Poseidon2CircuitRow = Poseidon2CircuitRow { + is_sponge: true, + reset: true, + absorb_flags: vec![true, false], + input_values: vec![ + Val::new(42), + Val::new(43), + Val::new(44), + Val::new(45), + Val::new(0), + Val::new(0), + Val::new(0), + Val::new(0), + ], + input_indices: vec![0; RATE_EXT], + output_indices: vec![0; RATE_EXT], + }; + + let trace = air.generate_trace_rows( + vec![sponge_a, sponge_b, sponge_c, sponge_d], + &constants, + fri_params.log_blowup, + perm, + ); type Dft = p3_dft::Radix2Bowers; let dft = Dft::default(); diff --git a/poseidon2-circuit-air/src/columns.rs b/poseidon2-circuit-air/src/columns.rs index 606ccb0..db3187c 100644 --- a/poseidon2-circuit-air/src/columns.rs +++ b/poseidon2-circuit-air/src/columns.rs @@ -21,7 +21,8 @@ use p3_poseidon2_air::Poseidon2Cols; pub struct Poseidon2CircuitCols< T, const WIDTH: usize, - const RATE: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -33,15 +34,15 @@ pub struct Poseidon2CircuitCols< pub is_sponge: T, pub reset: T, pub sponge_reset: T, - pub absorb_flags: [T; RATE], - pub input_indices: [T; WIDTH], - pub output_indices: [T; RATE], - pub _padding: [T; 290], + pub absorb_flags: [T; RATE_EXT], + pub input_indices: [T; WIDTH_EXT], + pub output_indices: [T; RATE_EXT], } pub const fn num_cols< const WIDTH: usize, - const RATE: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -51,7 +52,8 @@ pub const fn num_cols< Poseidon2CircuitCols< u8, WIDTH, - RATE, + WIDTH_EXT, + RATE_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -63,7 +65,8 @@ pub const fn num_cols< impl< T, const WIDTH: usize, - const RATE: usize, + const WIDTH_EXT: usize, + const RATE_EXT: usize, const SBOX_DEGREE: u64, const SBOX_REGISTERS: usize, const HALF_FULL_ROUNDS: usize, @@ -73,7 +76,8 @@ impl< Poseidon2CircuitCols< T, WIDTH, - RATE, + WIDTH_EXT, + RATE_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -86,7 +90,8 @@ impl< ) -> &Poseidon2CircuitCols< T, WIDTH, - RATE, + WIDTH_EXT, + RATE_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -96,7 +101,8 @@ impl< self.align_to:: &mut Poseidon2CircuitCols< T, WIDTH, - RATE, + WIDTH_EXT, + RATE_EXT, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, @@ -146,7 +155,8 @@ impl< self.align_to_mut::, T: Send + Sync + Clone> { + inner: M, + column_range: Range, + _phantom: core::marker::PhantomData, +} + +impl, T: Send + Sync + Clone> SubMatrixRowSlices { + /// Creates a new [`SubMatrixRowSlices`]. + #[must_use] + pub const fn new(inner: M, column_range: Range) -> Self { + Self { + inner, + column_range, + _phantom: core::marker::PhantomData, + } + } +} + +/// Implement `Matrix` for `SubMatrixRowSlices`. +impl, T: Send + Sync + Clone> Matrix for SubMatrixRowSlices { + #[inline] + fn row( + &self, + r: usize, + ) -> Option + Send + Sync>> { + self.inner.row(r).map(|row| { + row.into_iter() + .take(self.column_range.end) + .skip(self.column_range.start) + }) + } + + #[inline] + fn row_slice(&self, r: usize) -> Option> { + self.row(r)?.into_iter().collect::>().into() + } + + #[inline] + fn width(&self) -> usize { + self.column_range.len() + } + + #[inline] + fn height(&self) -> usize { + self.inner.height() + } +} + +/// A builder used to eval a sub-air. This will handle enforcing constraints for a subset of a +/// trace matrix. E.g. if a particular air needs to be enforced for a subset of the columns of +/// the trace, then the [`SubAirBuilder`] can be used. +pub struct SubAirBuilder<'a, AB: AirBuilder, SubAir: Air, T> { + inner: &'a mut AB, + column_range: Range, + _phantom: core::marker::PhantomData<(SubAir, T)>, +} + +impl<'a, AB: AirBuilder, SubAir: Air, T> SubAirBuilder<'a, AB, SubAir, T> { + /// Creates a new [`SubAirBuilder`]. + #[must_use] + pub fn new(inner: &'a mut AB, column_range: Range) -> Self { + Self { + inner, + column_range, + _phantom: core::marker::PhantomData, + } + } +} + +/// Implement `AirBuilder` for `SubAirBuilder`. +impl, F> AirBuilder for SubAirBuilder<'_, AB, SubAir, F> { + type F = AB::F; + type Expr = AB::Expr; + type Var = AB::Var; + type M = SubMatrixRowSlices; + + fn main(&self) -> Self::M { + let matrix = self.inner.main(); + + SubMatrixRowSlices::new(matrix, self.column_range.clone()) + } + + fn is_first_row(&self) -> Self::Expr { + self.inner.is_first_row() + } + + fn is_last_row(&self) -> Self::Expr { + self.inner.is_last_row() + } + + fn is_transition_window(&self, size: usize) -> Self::Expr { + self.inner.is_transition_window(size) + } + + fn assert_zero>(&mut self, x: I) { + self.inner.assert_zero(x.into()); + } +} From d949d5e44972c64d63573f7d35ef2434bacf17d4 Mon Sep 17 00:00:00 2001 From: hratoanina Date: Thu, 6 Nov 2025 10:07:02 +0100 Subject: [PATCH 9/9] Clippy --- poseidon2-circuit-air/src/air.rs | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/poseidon2-circuit-air/src/air.rs b/poseidon2-circuit-air/src/air.rs index 2bb82d0..4b5d653 100644 --- a/poseidon2-circuit-air/src/air.rs +++ b/poseidon2-circuit-air/src/air.rs @@ -135,27 +135,25 @@ impl< } let mut index_absorb = [false; RATE_EXT]; - for j in 0..RATE_EXT { - if absorb_flags[j] { - for k in 0..=j { - index_absorb[k] = true; + for (j, flag) in absorb_flags.iter().enumerate() { + if *flag { + for absorb in index_absorb.iter_mut().take(j + 1) { + *absorb = true; } } } - for j in 0..RATE_EXT { - if index_absorb[j] { + for (j, absorb) in index_absorb.iter_mut().enumerate() { + if *absorb { for d in 0..D { let idx = j * D + d; state[idx] = input_values[idx]; } - } else { - if *reset { - // During a reset, non-absorbed rate elements are zeroed. - for d in 0..D { - let idx = j * D + d; - state[idx] = F::ZERO; - } + } else if *reset { + // During a reset, non-absorbed rate elements are zeroed. + for d in 0..D { + let idx = j * D + d; + state[idx] = F::ZERO; } } } @@ -167,7 +165,7 @@ impl< } } - inputs.push(state.clone()); + inputs.push(state); state = perm.permute(state); } @@ -379,7 +377,7 @@ pub(crate) fn eval< .assert_zeros(arr); } - // let _is_squeeze = AB::Expr::ONE - current_absorb[0].clone(); + let _is_squeeze = AB::Expr::ONE - current_absorb[0].clone(); // TODO: Add all lookups: // - If current_absorb[i] = 1: // * local.rate[i] comes from input lookups.