From a5468becb47e5b3b55769dad0cd0775d96ac8e6b Mon Sep 17 00:00:00 2001 From: Henry de Valence Date: Tue, 19 Oct 2021 16:16:18 -0700 Subject: [PATCH] Split permutation from sponge construction Work towards #29; this doesn't touch the constraint system implementation yet, in order to be able to get design feedback on the software part. - The `poseidon::PoseidonParameters` struct is renamed to `poseidon::Parameters` but otherwise remains unchanged. - The `poseidon::PoseidonSpongeState` struct is renamed to `poseidon::State` and redefined to hold just the state itself, as well as the parameters needed to run the permutation. It exposes a `permute(&mut self)` method, `rate()` and `capacity()` accessors, as well as `Index`, `IndexMut`, `AsRef`, and `AsMut` impls that allow access to the state. - The `poseidon::PoseidonSponge` struct is renamed to `poseidon::Sponge` and holds a `State` and a `DuplexSpongeMode`. In other words, it consists of the state, together with the extra data tracking how that state is being used to implement a higher-level duplex construction. - The `CryptographicSponge` trait is changed so that `new()` takes an owned, `Self::Parameters`, not a borrowed one. This allows the caller to decide where to copy data, instead of forcing the sponge implementation to clone internally. Or, a `CryptographicSponge` implementation could declare the associated `Parameters` type to be some shared type (like an `Arc` wrapper) that avoids the need to copy at all. - The `SpongeExt` trait that allows converting back and forth between a state and a sponge is deleted; it's not safe to pass between abstraction layers that way. --- src/lib.rs | 18 +- src/poseidon/constraints.rs | 18 +- src/poseidon/mod.rs | 459 +----------------------------------- src/poseidon/parameters.rs | 54 +++++ src/poseidon/sponge.rs | 215 +++++++++++++++++ src/poseidon/state.rs | 135 +++++++++++ src/poseidon/tests.rs | 117 ++++++++- src/poseidon/traits.rs | 12 +- 8 files changed, 537 insertions(+), 491 deletions(-) create mode 100644 src/poseidon/parameters.rs create mode 100644 src/poseidon/sponge.rs create mode 100644 src/poseidon/state.rs diff --git a/src/lib.rs b/src/lib.rs index cfc0065..b1ae62f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,8 +28,8 @@ pub use absorb::*; /// The sponge for Poseidon /// -/// This implementation of Poseidon is entirely from Fractal's implementation in [COS20][cos] -/// with small syntax changes. +/// This implementation of Poseidon was derived from Fractal's implementation in +/// [COS20][cos]. /// /// [cos]: https://eprint.iacr.org/2019/1076 pub mod poseidon; @@ -110,7 +110,7 @@ pub trait CryptographicSponge: Clone { type Parameters; /// Initialize a new instance of the sponge. - fn new(params: &Self::Parameters) -> Self; + fn new(params: Self::Parameters) -> Self; /// Absorb an input into the sponge. fn absorb(&mut self, input: &impl Absorb); @@ -185,18 +185,6 @@ pub trait FieldBasedCryptographicSponge: CryptographicSponge { } } -/// An extension for the interface of a cryptographic sponge. -/// In addition to operations defined in `CryptographicSponge`, `SpongeExt` can convert itself to -/// a state, and instantiate itself from state. -pub trait SpongeExt: CryptographicSponge { - /// The full state of the cryptographic sponge. - type State: Clone; - /// Returns a sponge that uses `state`. - fn from_state(state: Self::State, params: &Self::Parameters) -> Self; - /// Consumes `self` and returns the state. - fn into_state(self) -> Self::State; -} - /// The mode structure for duplex sponges #[derive(Clone, Debug)] pub enum DuplexSpongeMode { diff --git a/src/poseidon/constraints.rs b/src/poseidon/constraints.rs index 50270cf..8b8edb5 100644 --- a/src/poseidon/constraints.rs +++ b/src/poseidon/constraints.rs @@ -1,6 +1,8 @@ +//! Constraints for Poseidon. + use crate::constraints::AbsorbGadget; use crate::constraints::{CryptographicSpongeVar, SpongeWithGadget}; -use crate::poseidon::{PoseidonParameters, PoseidonSponge}; +use crate::poseidon::{Parameters, Sponge}; use crate::DuplexSpongeMode; use ark_ff::{FpParameters, PrimeField}; use ark_r1cs_std::fields::fp::FpVar; @@ -21,7 +23,7 @@ pub struct PoseidonSpongeVar { pub cs: ConstraintSystemRef, /// Sponge Parameters - pub parameters: PoseidonParameters, + pub parameters: Parameters, // Sponge State /// The sponge's state @@ -30,7 +32,7 @@ pub struct PoseidonSpongeVar { pub mode: DuplexSpongeMode, } -impl SpongeWithGadget for PoseidonSponge { +impl SpongeWithGadget for Sponge { type Var = PoseidonSpongeVar; } @@ -179,11 +181,11 @@ impl PoseidonSpongeVar { } } -impl CryptographicSpongeVar> for PoseidonSpongeVar { - type Parameters = PoseidonParameters; +impl CryptographicSpongeVar> for PoseidonSpongeVar { + type Parameters = Parameters; #[tracing::instrument(target = "r1cs", skip(cs))] - fn new(cs: ConstraintSystemRef, parameters: &PoseidonParameters) -> Self { + fn new(cs: ConstraintSystemRef, parameters: &Parameters) -> Self { let zero = FpVar::::zero(); let state = vec![zero; parameters.rate + parameters.capacity]; let mode = DuplexSpongeMode::Absorbing { @@ -295,7 +297,7 @@ mod tests { use crate::constraints::CryptographicSpongeVar; use crate::poseidon::constraints::PoseidonSpongeVar; use crate::poseidon::tests::poseidon_parameters_for_test; - use crate::poseidon::PoseidonSponge; + use crate::poseidon::Sponge; use crate::{CryptographicSponge, FieldBasedCryptographicSponge}; use ark_ff::UniformRand; use ark_r1cs_std::fields::fp::FpVar; @@ -324,7 +326,7 @@ mod tests { let sponge_params = poseidon_parameters_for_test(); - let mut native_sponge = PoseidonSponge::::new(&sponge_params); + let mut native_sponge = Sponge::::new(sponge_params.clone()); let mut constraint_sponge = PoseidonSpongeVar::::new(cs.clone(), &sponge_params); native_sponge.absorb(&absorb1); diff --git a/src/poseidon/mod.rs b/src/poseidon/mod.rs index 506109b..76fbc48 100644 --- a/src/poseidon/mod.rs +++ b/src/poseidon/mod.rs @@ -1,459 +1,16 @@ -use crate::{ - batch_field_cast, squeeze_field_elements_with_sizes_default_impl, Absorb, CryptographicSponge, - DuplexSpongeMode, FieldBasedCryptographicSponge, FieldElementSize, SpongeExt, -}; -use ark_ff::{BigInteger, FpParameters, PrimeField}; -use ark_std::any::TypeId; -use ark_std::vec; -use ark_std::vec::Vec; - -/// constraints for Poseidon #[cfg(feature = "r1cs")] pub mod constraints; -#[cfg(test)] -mod tests; - -/// default parameters traits for Poseidon pub mod traits; -pub use traits::*; mod grain_lfsr; +mod parameters; +mod sponge; +mod state; -/// Parameters and RNG used -#[derive(Clone, Debug)] -pub struct PoseidonParameters { - /// Number of rounds in a full-round operation. - pub full_rounds: usize, - /// Number of rounds in a partial-round operation. - pub partial_rounds: usize, - /// Exponent used in S-boxes. - pub alpha: u64, - /// Additive Round keys. These are added before each MDS matrix application to make it an affine shift. - /// They are indexed by `ark[round_num][state_element_index]` - pub ark: Vec>, - /// Maximally Distance Separating (MDS) Matrix. - pub mds: Vec>, - /// The rate (in terms of number of field elements). - /// See [On the Indifferentiability of the Sponge Construction](https://iacr.org/archive/eurocrypt2008/49650180/49650180.pdf) - /// for more details on the rate and capacity of a sponge. - pub rate: usize, - /// The capacity (in terms of number of field elements). - pub capacity: usize, -} - -#[derive(Clone)] -/// A duplex sponge based using the Poseidon permutation. -/// -/// This implementation of Poseidon is entirely from Fractal's implementation in [COS20][cos] -/// with small syntax changes. -/// -/// [cos]: https://eprint.iacr.org/2019/1076 -pub struct PoseidonSponge { - /// Sponge Parameters - pub parameters: PoseidonParameters, - - // Sponge State - /// Current sponge's state (current elements in the permutation block) - pub state: Vec, - /// Current mode (whether its absorbing or squeezing) - pub mode: DuplexSpongeMode, -} - -impl PoseidonSponge { - fn apply_s_box(&self, state: &mut [F], is_full_round: bool) { - // Full rounds apply the S Box (x^alpha) to every element of state - if is_full_round { - for elem in state { - *elem = elem.pow(&[self.parameters.alpha]); - } - } - // Partial rounds apply the S Box (x^alpha) to just the first element of state - else { - state[0] = state[0].pow(&[self.parameters.alpha]); - } - } - - fn apply_ark(&self, state: &mut [F], round_number: usize) { - for (i, state_elem) in state.iter_mut().enumerate() { - state_elem.add_assign(&self.parameters.ark[round_number][i]); - } - } - - fn apply_mds(&self, state: &mut [F]) { - let mut new_state = Vec::new(); - for i in 0..state.len() { - let mut cur = F::zero(); - for (j, state_elem) in state.iter().enumerate() { - let term = state_elem.mul(&self.parameters.mds[i][j]); - cur.add_assign(&term); - } - new_state.push(cur); - } - state.clone_from_slice(&new_state[..state.len()]) - } - - fn permute(&mut self) { - let full_rounds_over_2 = self.parameters.full_rounds / 2; - let mut state = self.state.clone(); - for i in 0..full_rounds_over_2 { - self.apply_ark(&mut state, i); - self.apply_s_box(&mut state, true); - self.apply_mds(&mut state); - } - - for i in full_rounds_over_2..(full_rounds_over_2 + self.parameters.partial_rounds) { - self.apply_ark(&mut state, i); - self.apply_s_box(&mut state, false); - self.apply_mds(&mut state); - } - - for i in (full_rounds_over_2 + self.parameters.partial_rounds) - ..(self.parameters.partial_rounds + self.parameters.full_rounds) - { - self.apply_ark(&mut state, i); - self.apply_s_box(&mut state, true); - self.apply_mds(&mut state); - } - self.state = state; - } - - // Absorbs everything in elements, this does not end in an absorbtion. - fn absorb_internal(&mut self, mut rate_start_index: usize, elements: &[F]) { - let mut remaining_elements = elements; - - loop { - // if we can finish in this call - if rate_start_index + remaining_elements.len() <= self.parameters.rate { - for (i, element) in remaining_elements.iter().enumerate() { - self.state[self.parameters.capacity + i + rate_start_index] += element; - } - self.mode = DuplexSpongeMode::Absorbing { - next_absorb_index: rate_start_index + remaining_elements.len(), - }; - - return; - } - // otherwise absorb (rate - rate_start_index) elements - let num_elements_absorbed = self.parameters.rate - rate_start_index; - for (i, element) in remaining_elements - .iter() - .enumerate() - .take(num_elements_absorbed) - { - self.state[self.parameters.capacity + i + rate_start_index] += element; - } - self.permute(); - // the input elements got truncated by num elements absorbed - remaining_elements = &remaining_elements[num_elements_absorbed..]; - rate_start_index = 0; - } - } - - // Squeeze |output| many elements. This does not end in a squeeze - fn squeeze_internal(&mut self, mut rate_start_index: usize, output: &mut [F]) { - let mut output_remaining = output; - loop { - // if we can finish in this call - if rate_start_index + output_remaining.len() <= self.parameters.rate { - output_remaining.clone_from_slice( - &self.state[self.parameters.capacity + rate_start_index - ..(self.parameters.capacity + output_remaining.len() + rate_start_index)], - ); - self.mode = DuplexSpongeMode::Squeezing { - next_squeeze_index: rate_start_index + output_remaining.len(), - }; - return; - } - // otherwise squeeze (rate - rate_start_index) elements - let num_elements_squeezed = self.parameters.rate - rate_start_index; - output_remaining[..num_elements_squeezed].clone_from_slice( - &self.state[self.parameters.capacity + rate_start_index - ..(self.parameters.capacity + num_elements_squeezed + rate_start_index)], - ); - - // Unless we are done with squeezing in this call, permute. - if output_remaining.len() != self.parameters.rate { - self.permute(); - } - // Repeat with updated output slices - output_remaining = &mut output_remaining[num_elements_squeezed..]; - rate_start_index = 0; - } - } -} - -impl PoseidonParameters { - /// Initialize the parameter for Poseidon Sponge. - pub fn new( - full_rounds: usize, - partial_rounds: usize, - alpha: u64, - mds: Vec>, - ark: Vec>, - rate: usize, - capacity: usize, - ) -> Self { - assert_eq!(ark.len(), full_rounds + partial_rounds); - for item in &ark { - assert_eq!(item.len(), rate + capacity); - } - assert_eq!(mds.len(), rate + capacity); - for item in &mds { - assert_eq!(item.len(), rate + capacity); - } - Self { - full_rounds, - partial_rounds, - alpha, - mds, - ark, - rate, - capacity, - } - } -} - -impl CryptographicSponge for PoseidonSponge { - type Parameters = PoseidonParameters; - - fn new(parameters: &Self::Parameters) -> Self { - let state = vec![F::zero(); parameters.rate + parameters.capacity]; - let mode = DuplexSpongeMode::Absorbing { - next_absorb_index: 0, - }; - - Self { - parameters: parameters.clone(), - state, - mode, - } - } - - fn absorb(&mut self, input: &impl Absorb) { - let elems = input.to_sponge_field_elements_as_vec::(); - if elems.is_empty() { - return; - } - - match self.mode { - DuplexSpongeMode::Absorbing { next_absorb_index } => { - let mut absorb_index = next_absorb_index; - if absorb_index == self.parameters.rate { - self.permute(); - absorb_index = 0; - } - self.absorb_internal(absorb_index, elems.as_slice()); - } - DuplexSpongeMode::Squeezing { - next_squeeze_index: _, - } => { - self.permute(); - self.absorb_internal(0, elems.as_slice()); - } - }; - } - - fn squeeze_bytes(&mut self, num_bytes: usize) -> Vec { - let usable_bytes = (F::Params::CAPACITY / 8) as usize; - - let num_elements = (num_bytes + usable_bytes - 1) / usable_bytes; - let src_elements = self.squeeze_native_field_elements(num_elements); - - let mut bytes: Vec = Vec::with_capacity(usable_bytes * num_elements); - for elem in &src_elements { - let elem_bytes = elem.into_repr().to_bytes_le(); - bytes.extend_from_slice(&elem_bytes[..usable_bytes]); - } - - bytes.truncate(num_bytes); - bytes - } - - fn squeeze_bits(&mut self, num_bits: usize) -> Vec { - let usable_bits = F::Params::CAPACITY as usize; - - let num_elements = (num_bits + usable_bits - 1) / usable_bits; - let src_elements = self.squeeze_native_field_elements(num_elements); - - let mut bits: Vec = Vec::with_capacity(usable_bits * num_elements); - for elem in &src_elements { - let elem_bits = elem.into_repr().to_bits_le(); - bits.extend_from_slice(&elem_bits[..usable_bits]); - } - - bits.truncate(num_bits); - bits - } - - fn squeeze_field_elements_with_sizes( - &mut self, - sizes: &[FieldElementSize], - ) -> Vec { - if F::characteristic() == F2::characteristic() { - // native case - let mut buf = Vec::with_capacity(sizes.len()); - batch_field_cast( - &self.squeeze_native_field_elements_with_sizes(sizes), - &mut buf, - ) - .unwrap(); - buf - } else { - squeeze_field_elements_with_sizes_default_impl(self, sizes) - } - } - - fn squeeze_field_elements(&mut self, num_elements: usize) -> Vec { - if TypeId::of::() == TypeId::of::() { - let result = self.squeeze_native_field_elements(num_elements); - let mut cast = Vec::with_capacity(result.len()); - batch_field_cast(&result, &mut cast).unwrap(); - cast - } else { - self.squeeze_field_elements_with_sizes::( - vec![FieldElementSize::Full; num_elements].as_slice(), - ) - } - } -} - -impl FieldBasedCryptographicSponge for PoseidonSponge { - fn squeeze_native_field_elements(&mut self, num_elements: usize) -> Vec { - let mut squeezed_elems = vec![F::zero(); num_elements]; - match self.mode { - DuplexSpongeMode::Absorbing { - next_absorb_index: _, - } => { - self.permute(); - self.squeeze_internal(0, &mut squeezed_elems); - } - DuplexSpongeMode::Squeezing { next_squeeze_index } => { - let mut squeeze_index = next_squeeze_index; - if squeeze_index == self.parameters.rate { - self.permute(); - squeeze_index = 0; - } - self.squeeze_internal(squeeze_index, &mut squeezed_elems); - } - }; - - squeezed_elems - } -} - -#[derive(Clone)] -/// Stores the state of a Poseidon Sponge. Does not store any parameter. -pub struct PoseidonSpongeState { - state: Vec, - mode: DuplexSpongeMode, -} - -impl SpongeExt for PoseidonSponge { - type State = PoseidonSpongeState; - - fn from_state(state: Self::State, params: &Self::Parameters) -> Self { - let mut sponge = Self::new(params); - sponge.mode = state.mode; - sponge.state = state.state; - sponge - } - - fn into_state(self) -> Self::State { - Self::State { - state: self.state, - mode: self.mode, - } - } -} +pub use parameters::Parameters; +pub use sponge::Sponge; +pub use state::State; +pub use traits::*; #[cfg(test)] -mod test { - use crate::poseidon::{ - PoseidonDefaultParameters, PoseidonDefaultParametersEntry, PoseidonDefaultParametersField, - }; - use crate::{poseidon::PoseidonSponge, CryptographicSponge, FieldBasedCryptographicSponge}; - use ark_ff::{field_new, BigInteger256, FftParameters, Fp256, Fp256Parameters, FpParameters}; - use ark_test_curves::bls12_381::FrParameters; - - pub struct TestFrParameters; - - impl Fp256Parameters for TestFrParameters {} - impl FftParameters for TestFrParameters { - type BigInt = ::BigInt; - const TWO_ADICITY: u32 = FrParameters::TWO_ADICITY; - const TWO_ADIC_ROOT_OF_UNITY: Self::BigInt = FrParameters::TWO_ADIC_ROOT_OF_UNITY; - } - - // This TestFrParameters is the same as the BLS12-381's Fr. - // MODULUS = 52435875175126190479447740508185965837690552500527637822603658699938581184513 - impl FpParameters for TestFrParameters { - const MODULUS: BigInteger256 = FrParameters::MODULUS; - const MODULUS_BITS: u32 = FrParameters::MODULUS_BITS; - const CAPACITY: u32 = FrParameters::CAPACITY; - const REPR_SHAVE_BITS: u32 = FrParameters::REPR_SHAVE_BITS; - const R: BigInteger256 = FrParameters::R; - const R2: BigInteger256 = FrParameters::R2; - const INV: u64 = FrParameters::INV; - const GENERATOR: BigInteger256 = FrParameters::GENERATOR; - const MODULUS_MINUS_ONE_DIV_TWO: BigInteger256 = FrParameters::MODULUS_MINUS_ONE_DIV_TWO; - const T: BigInteger256 = FrParameters::T; - const T_MINUS_ONE_DIV_TWO: BigInteger256 = FrParameters::T_MINUS_ONE_DIV_TWO; - } - - impl PoseidonDefaultParameters for TestFrParameters { - const PARAMS_OPT_FOR_CONSTRAINTS: [PoseidonDefaultParametersEntry; 7] = [ - PoseidonDefaultParametersEntry::new(2, 17, 8, 31, 0), - PoseidonDefaultParametersEntry::new(3, 5, 8, 56, 0), - PoseidonDefaultParametersEntry::new(4, 5, 8, 56, 0), - PoseidonDefaultParametersEntry::new(5, 5, 8, 57, 0), - PoseidonDefaultParametersEntry::new(6, 5, 8, 57, 0), - PoseidonDefaultParametersEntry::new(7, 5, 8, 57, 0), - PoseidonDefaultParametersEntry::new(8, 5, 8, 57, 0), - ]; - const PARAMS_OPT_FOR_WEIGHTS: [PoseidonDefaultParametersEntry; 7] = [ - PoseidonDefaultParametersEntry::new(2, 257, 8, 13, 0), - PoseidonDefaultParametersEntry::new(3, 257, 8, 13, 0), - PoseidonDefaultParametersEntry::new(4, 257, 8, 13, 0), - PoseidonDefaultParametersEntry::new(5, 257, 8, 13, 0), - PoseidonDefaultParametersEntry::new(6, 257, 8, 13, 0), - PoseidonDefaultParametersEntry::new(7, 257, 8, 13, 0), - PoseidonDefaultParametersEntry::new(8, 257, 8, 13, 0), - ]; - } - - pub type TestFr = Fp256; - - #[test] - fn test_poseidon_sponge_consistency() { - let sponge_param = TestFr::get_default_poseidon_parameters(2, false).unwrap(); - - let mut sponge = PoseidonSponge::::new(&sponge_param); - sponge.absorb(&vec![ - TestFr::from(0u8), - TestFr::from(1u8), - TestFr::from(2u8), - ]); - let res = sponge.squeeze_native_field_elements(3); - assert_eq!( - res[0], - field_new!( - TestFr, - "40442793463571304028337753002242186710310163897048962278675457993207843616876" - ) - ); - assert_eq!( - res[1], - field_new!( - TestFr, - "2664374461699898000291153145224099287711224021716202960480903840045233645301" - ) - ); - assert_eq!( - res[2], - field_new!( - TestFr, - "50191078828066923662070228256530692951801504043422844038937334196346054068797" - ) - ); - } -} +mod tests; diff --git a/src/poseidon/parameters.rs b/src/poseidon/parameters.rs new file mode 100644 index 0000000..472abde --- /dev/null +++ b/src/poseidon/parameters.rs @@ -0,0 +1,54 @@ +use ark_ff::PrimeField; + +/// Parameters describing a Poseidon instance. +#[derive(Clone, Debug)] +pub struct Parameters { + /// Number of rounds in a full-round operation. + pub full_rounds: usize, + /// Number of rounds in a partial-round operation. + pub partial_rounds: usize, + /// Exponent used in S-boxes. + pub alpha: u64, + /// Additive Round keys. These are added before each MDS matrix application to make it an affine shift. + /// They are indexed by `ark[round_num][state_element_index]` + pub ark: Vec>, + /// Maximally Distance Separating (MDS) Matrix. + pub mds: Vec>, + /// The rate (in terms of number of field elements). + /// See [On the Indifferentiability of the Sponge Construction](https://iacr.org/archive/eurocrypt2008/49650180/49650180.pdf) + /// for more details on the rate and capacity of a sponge. + pub rate: usize, + /// The capacity (in terms of number of field elements). + pub capacity: usize, +} + +impl Parameters { + /// Initialize the parameter for Poseidon Sponge. + pub fn new( + full_rounds: usize, + partial_rounds: usize, + alpha: u64, + mds: Vec>, + ark: Vec>, + rate: usize, + capacity: usize, + ) -> Self { + assert_eq!(ark.len(), full_rounds + partial_rounds); + for item in &ark { + assert_eq!(item.len(), rate + capacity); + } + assert_eq!(mds.len(), rate + capacity); + for item in &mds { + assert_eq!(item.len(), rate + capacity); + } + Self { + full_rounds, + partial_rounds, + alpha, + mds, + ark, + rate, + capacity, + } + } +} diff --git a/src/poseidon/sponge.rs b/src/poseidon/sponge.rs new file mode 100644 index 0000000..e6a4b6e --- /dev/null +++ b/src/poseidon/sponge.rs @@ -0,0 +1,215 @@ +use ark_ff::{BigInteger, FpParameters, PrimeField}; +use ark_std::{any::TypeId, vec::Vec}; + +use crate::{ + batch_field_cast, squeeze_field_elements_with_sizes_default_impl, Absorb, CryptographicSponge, + DuplexSpongeMode, FieldBasedCryptographicSponge, FieldElementSize, +}; + +use super::{Parameters, State}; + +/// A duplex sponge based using the Poseidon permutation. +/// +/// This implementation of Poseidon was derived from Fractal's implementation in +/// [COS20][cos]. +/// +/// [cos]: https://eprint.iacr.org/2019/1076 +#[derive(Clone)] +pub struct Sponge { + /// The underlying sponge state. + pub state: State, + /// Current mode (whether its absorbing or squeezing) + pub mode: DuplexSpongeMode, +} + +impl Sponge { + // Absorbs everything in elements, this does not end in an absorbtion. + fn absorb_internal(&mut self, mut rate_start_index: usize, elements: &[F]) { + let (rate, capacity) = (self.state.rate(), self.state.capacity()); + let mut remaining_elements = elements; + + loop { + // if we can finish in this call + if rate_start_index + remaining_elements.len() <= rate { + for (i, element) in remaining_elements.iter().enumerate() { + self.state[capacity + i + rate_start_index] += element; + } + self.mode = DuplexSpongeMode::Absorbing { + next_absorb_index: rate_start_index + remaining_elements.len(), + }; + + return; + } + // otherwise absorb (rate - rate_start_index) elements + let num_elements_absorbed = rate - rate_start_index; + for (i, element) in remaining_elements + .iter() + .enumerate() + .take(num_elements_absorbed) + { + self.state[capacity + i + rate_start_index] += element; + } + self.state.permute(); + // the input elements got truncated by num elements absorbed + remaining_elements = &remaining_elements[num_elements_absorbed..]; + rate_start_index = 0; + } + } + + // Squeeze |output| many elements. This does not end in a squeeze + fn squeeze_internal(&mut self, mut rate_start_index: usize, output: &mut [F]) { + let (rate, capacity) = (self.state.rate(), self.state.capacity()); + let mut output_remaining = output; + loop { + // if we can finish in this call + if rate_start_index + output_remaining.len() <= rate { + output_remaining.clone_from_slice( + &self.state[capacity + rate_start_index + ..(capacity + output_remaining.len() + rate_start_index)], + ); + self.mode = DuplexSpongeMode::Squeezing { + next_squeeze_index: rate_start_index + output_remaining.len(), + }; + return; + } + // otherwise squeeze (rate - rate_start_index) elements + let num_elements_squeezed = rate - rate_start_index; + output_remaining[..num_elements_squeezed].clone_from_slice( + &self.state[capacity + rate_start_index + ..(capacity + num_elements_squeezed + rate_start_index)], + ); + + // Unless we are done with squeezing in this call, permute. + if output_remaining.len() != rate { + self.state.permute(); + } + // Repeat with updated output slices + output_remaining = &mut output_remaining[num_elements_squeezed..]; + rate_start_index = 0; + } + } +} + +impl CryptographicSponge for Sponge { + type Parameters = Parameters; + + fn new(parameters: Self::Parameters) -> Self { + Self { + state: parameters.into(), + mode: DuplexSpongeMode::Absorbing { + next_absorb_index: 0, + }, + } + } + + fn absorb(&mut self, input: &impl Absorb) { + let elems = input.to_sponge_field_elements_as_vec::(); + if elems.is_empty() { + return; + } + + match self.mode { + DuplexSpongeMode::Absorbing { next_absorb_index } => { + let mut absorb_index = next_absorb_index; + if absorb_index == self.state.rate() { + self.state.permute(); + absorb_index = 0; + } + self.absorb_internal(absorb_index, elems.as_slice()); + } + DuplexSpongeMode::Squeezing { + next_squeeze_index: _, + } => { + self.state.permute(); + self.absorb_internal(0, elems.as_slice()); + } + }; + } + + fn squeeze_bytes(&mut self, num_bytes: usize) -> Vec { + let usable_bytes = (F::Params::CAPACITY / 8) as usize; + + let num_elements = (num_bytes + usable_bytes - 1) / usable_bytes; + let src_elements = self.squeeze_native_field_elements(num_elements); + + let mut bytes: Vec = Vec::with_capacity(usable_bytes * num_elements); + for elem in &src_elements { + let elem_bytes = elem.into_repr().to_bytes_le(); + bytes.extend_from_slice(&elem_bytes[..usable_bytes]); + } + + bytes.truncate(num_bytes); + bytes + } + + fn squeeze_bits(&mut self, num_bits: usize) -> Vec { + let usable_bits = F::Params::CAPACITY as usize; + + let num_elements = (num_bits + usable_bits - 1) / usable_bits; + let src_elements = self.squeeze_native_field_elements(num_elements); + + let mut bits: Vec = Vec::with_capacity(usable_bits * num_elements); + for elem in &src_elements { + let elem_bits = elem.into_repr().to_bits_le(); + bits.extend_from_slice(&elem_bits[..usable_bits]); + } + + bits.truncate(num_bits); + bits + } + + fn squeeze_field_elements_with_sizes( + &mut self, + sizes: &[FieldElementSize], + ) -> Vec { + if F::characteristic() == F2::characteristic() { + // native case + let mut buf = Vec::with_capacity(sizes.len()); + batch_field_cast( + &self.squeeze_native_field_elements_with_sizes(sizes), + &mut buf, + ) + .unwrap(); + buf + } else { + squeeze_field_elements_with_sizes_default_impl(self, sizes) + } + } + + fn squeeze_field_elements(&mut self, num_elements: usize) -> Vec { + if TypeId::of::() == TypeId::of::() { + let result = self.squeeze_native_field_elements(num_elements); + let mut cast = Vec::with_capacity(result.len()); + batch_field_cast(&result, &mut cast).unwrap(); + cast + } else { + self.squeeze_field_elements_with_sizes::( + vec![FieldElementSize::Full; num_elements].as_slice(), + ) + } + } +} + +impl FieldBasedCryptographicSponge for Sponge { + fn squeeze_native_field_elements(&mut self, num_elements: usize) -> Vec { + let mut squeezed_elems = vec![F::zero(); num_elements]; + match self.mode { + DuplexSpongeMode::Absorbing { + next_absorb_index: _, + } => { + self.state.permute(); + self.squeeze_internal(0, &mut squeezed_elems); + } + DuplexSpongeMode::Squeezing { next_squeeze_index } => { + let mut squeeze_index = next_squeeze_index; + if squeeze_index == self.state.rate() { + self.state.permute(); + squeeze_index = 0; + } + self.squeeze_internal(squeeze_index, &mut squeezed_elems); + } + }; + + squeezed_elems + } +} diff --git a/src/poseidon/state.rs b/src/poseidon/state.rs new file mode 100644 index 0000000..41edb46 --- /dev/null +++ b/src/poseidon/state.rs @@ -0,0 +1,135 @@ +use ark_ff::PrimeField; +use std::{ + ops::{Index, IndexMut}, + slice::SliceIndex, +}; + +use super::Parameters; + +#[derive(Clone)] +/// A raw Poseidon state, with direct access to the permutation. +/// +/// This is a lower-level API than +/// [`poseidon::Sponge`](crate::poseidon::Sponge), which builds a duplex +/// construction. +/// +/// A new state is constructed from a [`Parameters`] instance using [`From`] or +/// [`Into`]. The state itself can be accessed and updated using [`Index`] and +/// [`IndexMut`] or [`AsRef`] and [`AsMut`]. +pub struct State { + // Use a Box<[F]> rather than Vec because the state is not resizable. + state: Box<[F]>, + parameters: Parameters, +} + +impl From> for State { + fn from(parameters: Parameters) -> Self { + let state = vec![F::zero(); parameters.capacity + parameters.rate].into_boxed_slice(); + Self { state, parameters } + } +} + +impl Index for State +where + F: PrimeField, + I: SliceIndex<[F]>, +{ + type Output = >::Output; + #[inline] + fn index(&self, index: I) -> &Self::Output { + self.state.index(index) + } +} + +impl IndexMut for State +where + F: PrimeField, + I: SliceIndex<[F]>, +{ + #[inline] + fn index_mut(&mut self, index: I) -> &mut Self::Output { + self.state.index_mut(index) + } +} + +impl AsRef<[F]> for State { + #[inline] + fn as_ref(&self) -> &[F] { + self.state.as_ref() + } +} + +impl AsMut<[F]> for State { + #[inline] + fn as_mut(&mut self) -> &mut [F] { + self.state.as_mut() + } +} + +impl State { + fn apply_s_box(&mut self, is_full_round: bool) { + // Full rounds apply the S Box (x^alpha) to every element of state + if is_full_round { + for elem in self.state.iter_mut() { + *elem = elem.pow(&[self.parameters.alpha]); + } + } + // Partial rounds apply the S Box (x^alpha) to just the first element of state + else { + self.state[0] = self.state[0].pow(&[self.parameters.alpha]); + } + } + + fn apply_ark(&mut self, round_number: usize) { + for (i, state_elem) in self.state.iter_mut().enumerate() { + state_elem.add_assign(&self.parameters.ark[round_number][i]); + } + } + + fn apply_mds(&mut self) { + let mut new_state = Vec::with_capacity(self.state.len()); + for i in 0..self.state.len() { + let mut cur = F::zero(); + for (j, state_elem) in self.state.iter().enumerate() { + let term = state_elem.mul(&self.parameters.mds[i][j]); + cur.add_assign(&term); + } + new_state.push(cur); + } + self.state = new_state.into_boxed_slice(); + } + + /// Runs the permutation, updating the state. + pub fn permute(&mut self) { + let full_rounds_over_2 = self.parameters.full_rounds / 2; + for i in 0..full_rounds_over_2 { + self.apply_ark(i); + self.apply_s_box(true); + self.apply_mds(); + } + + for i in full_rounds_over_2..(full_rounds_over_2 + self.parameters.partial_rounds) { + self.apply_ark(i); + self.apply_s_box(false); + self.apply_mds(); + } + + for i in (full_rounds_over_2 + self.parameters.partial_rounds) + ..(self.parameters.partial_rounds + self.parameters.full_rounds) + { + self.apply_ark(i); + self.apply_s_box(true); + self.apply_mds(); + } + } + + /// Returns the rate of the permutation. + pub fn rate(&self) -> usize { + self.parameters.rate + } + + /// Returns the capacity of the permutation. + pub fn capacity(&self) -> usize { + self.parameters.capacity + } +} diff --git a/src/poseidon/tests.rs b/src/poseidon/tests.rs index 62429e9..a0d6857 100644 --- a/src/poseidon/tests.rs +++ b/src/poseidon/tests.rs @@ -1,17 +1,110 @@ -use crate::poseidon::{PoseidonParameters, PoseidonSponge}; -use crate::{absorb, collect_sponge_bytes, collect_sponge_field_elements}; -use crate::{Absorb, AbsorbWithLength, CryptographicSponge, FieldBasedCryptographicSponge}; -use ark_ff::{One, PrimeField, UniformRand}; +use ark_ff::{ + field_new, BigInteger256, FftParameters, Fp256, Fp256Parameters, FpParameters, One, PrimeField, + UniformRand, +}; use ark_std::test_rng; -use ark_test_curves::bls12_381::Fr; +use ark_test_curves::bls12_381::{Fr, FrParameters}; + +use crate::{ + absorb, collect_sponge_bytes, collect_sponge_field_elements, Absorb, AbsorbWithLength, + CryptographicSponge, FieldBasedCryptographicSponge, +}; + +use super::{ + Parameters, PoseidonDefaultParameters, PoseidonDefaultParametersEntry, + PoseidonDefaultParametersField, Sponge, +}; + +pub struct TestFrParameters; + +impl Fp256Parameters for TestFrParameters {} +impl FftParameters for TestFrParameters { + type BigInt = ::BigInt; + const TWO_ADICITY: u32 = FrParameters::TWO_ADICITY; + const TWO_ADIC_ROOT_OF_UNITY: Self::BigInt = FrParameters::TWO_ADIC_ROOT_OF_UNITY; +} + +// This TestFrParameters is the same as the BLS12-381's Fr. +// MODULUS = 52435875175126190479447740508185965837690552500527637822603658699938581184513 +impl FpParameters for TestFrParameters { + const MODULUS: BigInteger256 = FrParameters::MODULUS; + const MODULUS_BITS: u32 = FrParameters::MODULUS_BITS; + const CAPACITY: u32 = FrParameters::CAPACITY; + const REPR_SHAVE_BITS: u32 = FrParameters::REPR_SHAVE_BITS; + const R: BigInteger256 = FrParameters::R; + const R2: BigInteger256 = FrParameters::R2; + const INV: u64 = FrParameters::INV; + const GENERATOR: BigInteger256 = FrParameters::GENERATOR; + const MODULUS_MINUS_ONE_DIV_TWO: BigInteger256 = FrParameters::MODULUS_MINUS_ONE_DIV_TWO; + const T: BigInteger256 = FrParameters::T; + const T_MINUS_ONE_DIV_TWO: BigInteger256 = FrParameters::T_MINUS_ONE_DIV_TWO; +} + +impl PoseidonDefaultParameters for TestFrParameters { + const PARAMS_OPT_FOR_CONSTRAINTS: [PoseidonDefaultParametersEntry; 7] = [ + PoseidonDefaultParametersEntry::new(2, 17, 8, 31, 0), + PoseidonDefaultParametersEntry::new(3, 5, 8, 56, 0), + PoseidonDefaultParametersEntry::new(4, 5, 8, 56, 0), + PoseidonDefaultParametersEntry::new(5, 5, 8, 57, 0), + PoseidonDefaultParametersEntry::new(6, 5, 8, 57, 0), + PoseidonDefaultParametersEntry::new(7, 5, 8, 57, 0), + PoseidonDefaultParametersEntry::new(8, 5, 8, 57, 0), + ]; + const PARAMS_OPT_FOR_WEIGHTS: [PoseidonDefaultParametersEntry; 7] = [ + PoseidonDefaultParametersEntry::new(2, 257, 8, 13, 0), + PoseidonDefaultParametersEntry::new(3, 257, 8, 13, 0), + PoseidonDefaultParametersEntry::new(4, 257, 8, 13, 0), + PoseidonDefaultParametersEntry::new(5, 257, 8, 13, 0), + PoseidonDefaultParametersEntry::new(6, 257, 8, 13, 0), + PoseidonDefaultParametersEntry::new(7, 257, 8, 13, 0), + PoseidonDefaultParametersEntry::new(8, 257, 8, 13, 0), + ]; +} + +pub type TestFr = Fp256; + +#[test] +fn test_poseidon_sponge_consistency() { + let sponge_param = TestFr::get_default_poseidon_parameters(2, false).unwrap(); + + let mut sponge = Sponge::::new(sponge_param); + sponge.absorb(&vec![ + TestFr::from(0u8), + TestFr::from(1u8), + TestFr::from(2u8), + ]); + let res = sponge.squeeze_native_field_elements(3); + assert_eq!( + res[0], + field_new!( + TestFr, + "40442793463571304028337753002242186710310163897048962278675457993207843616876" + ) + ); + assert_eq!( + res[1], + field_new!( + TestFr, + "2664374461699898000291153145224099287711224021716202960480903840045233645301" + ) + ); + assert_eq!( + res[2], + field_new!( + TestFr, + "50191078828066923662070228256530692951801504043422844038937334196346054068797" + ) + ); +} + fn assert_different_encodings(a: &A, b: &A) { let bytes1 = a.to_sponge_bytes_as_vec(); let bytes2 = b.to_sponge_bytes_as_vec(); assert_ne!(bytes1, bytes2); let sponge_param = poseidon_parameters_for_test(); - let mut sponge1 = PoseidonSponge::::new(&sponge_param); - let mut sponge2 = PoseidonSponge::::new(&sponge_param); + let mut sponge1 = Sponge::::new(sponge_param.clone()); + let mut sponge2 = Sponge::::new(sponge_param); sponge1.absorb(&a); sponge2.absorb(&b); @@ -72,7 +165,7 @@ fn test_squeeze_cast_native() { let mut rng = test_rng(); let sponge_param = poseidon_parameters_for_test(); let elem = Fr::rand(&mut rng); - let mut sponge1 = PoseidonSponge::::new(&sponge_param); + let mut sponge1 = Sponge::::new(sponge_param); sponge1.absorb(&elem); let mut sponge2 = sponge1.clone(); @@ -86,11 +179,11 @@ fn test_squeeze_cast_native() { #[test] fn test_macros() { let sponge_param = poseidon_parameters_for_test(); - let mut sponge1 = PoseidonSponge::::new(&sponge_param); + let mut sponge1 = Sponge::::new(sponge_param.clone()); sponge1.absorb(&vec![1, 2, 3, 4, 5, 6]); sponge1.absorb(&Fr::from(114514u128)); - let mut sponge2 = PoseidonSponge::::new(&sponge_param); + let mut sponge2 = Sponge::::new(sponge_param); absorb!(&mut sponge2, vec![1, 2, 3, 4, 5, 6], Fr::from(114514u128)); let expected = sponge1.squeeze_native_field_elements(3); @@ -116,7 +209,7 @@ fn test_macros() { } /// Generate default parameters (bls381-fr-only) for alpha = 17, state-size = 8 -pub(crate) fn poseidon_parameters_for_test() -> PoseidonParameters { +pub(crate) fn poseidon_parameters_for_test() -> Parameters { let alpha = 17; let mds = vec![ vec![ @@ -807,7 +900,7 @@ pub(crate) fn poseidon_parameters_for_test() -> PoseidonParameter let partial_rounds = total_rounds - full_rounds; let capacity = 1; let rate = 2; - PoseidonParameters { + Parameters { full_rounds, partial_rounds, alpha, diff --git a/src/poseidon/traits.rs b/src/poseidon/traits.rs index d288be9..1a58bb7 100644 --- a/src/poseidon/traits.rs +++ b/src/poseidon/traits.rs @@ -1,5 +1,7 @@ +//! default parameters traits for Poseidon + use crate::poseidon::grain_lfsr::PoseidonGrainLFSR; -use crate::poseidon::PoseidonParameters; +use crate::poseidon::Parameters; use ark_ff::{fields::models::*, FpParameters, PrimeField}; use ark_std::{vec, vec::Vec}; @@ -62,14 +64,14 @@ pub trait PoseidonDefaultParametersField: PrimeField { fn get_default_poseidon_parameters( rate: usize, optimized_for_weights: bool, - ) -> Option>; + ) -> Option>; } /// Internal function that uses the `PoseidonDefaultParameters` to compute the Poseidon parameters. pub fn get_default_poseidon_parameters_internal( rate: usize, optimized_for_weights: bool, -) -> Option> { +) -> Option> { let params_set = if !optimized_for_weights { P::PARAMS_OPT_FOR_CONSTRAINTS } else { @@ -86,7 +88,7 @@ pub fn get_default_poseidon_parameters_internal Option> { + ) -> Option> { get_default_poseidon_parameters_internal::(rate, optimized_for_weights) } }