From 423006b905b208c034091c086ebbc5093af874d6 Mon Sep 17 00:00:00 2001 From: Jack Grigg Date: Fri, 10 Dec 2021 02:25:09 +0000 Subject: [PATCH] poseidon: Remove redundant additions when squeezing sponge In the previous commit, we fixed a bug where padding was being added to the state when the sponge was in squeezing mode. But there's no need to assign a circuit region in which we add constant zeroes to the state :) --- src/circuit/gadget/poseidon.rs | 32 ++++++++++++++--------------- src/circuit/gadget/poseidon/pow5.rs | 27 ++++++++++++++---------- src/primitives/poseidon.rs | 30 ++++++++++++++------------- 3 files changed, 47 insertions(+), 42 deletions(-) diff --git a/src/circuit/gadget/poseidon.rs b/src/circuit/gadget/poseidon.rs index aff0c2c5..d851aacf 100644 --- a/src/circuit/gadget/poseidon.rs +++ b/src/circuit/gadget/poseidon.rs @@ -16,7 +16,7 @@ mod pow5; pub use pow5::{Pow5Chip, Pow5Config, StateWord}; use crate::primitives::poseidon::{ - Absorbing, ConstantLength, Domain, Spec, SpongeMode, SpongeRate, Squeezing, State, + Absorbing, ConstantLength, Domain, Spec, SpongeMode, Squeezing, State, }; /// A word from the padded input to a Poseidon sponge. @@ -63,11 +63,11 @@ pub trait PoseidonSpongeInstructions< &self, layouter: &mut impl Layouter, initial_state: &State, - input: &SpongeRate, RATE>, + input: &Absorbing, RATE>, ) -> Result, Error>; /// Extracts sponge output from the given state. - fn get_output(state: &State) -> SpongeRate; + fn get_output(state: &State) -> Squeezing; } /// A word over which the Poseidon permutation operates. @@ -112,9 +112,11 @@ fn poseidon_sponge< chip: &PoseidonChip, mut layouter: impl Layouter, state: &mut State, - input: &SpongeRate, RATE>, -) -> Result, Error> { - *state = chip.add_input(&mut layouter, state, input)?; + input: Option<&Absorbing, RATE>>, +) -> Result, Error> { + if let Some(input) = input { + *state = chip.add_input(&mut layouter, state, input)?; + } *state = chip.permute(&mut layouter, state)?; Ok(PoseidonChip::get_output(state)) } @@ -179,7 +181,7 @@ impl< &self.chip, layouter.namespace(|| "PoseidonSponge"), &mut self.state, - &self.mode.0, + Some(&self.mode), )?; self.mode = Absorbing::init_with(value); @@ -193,12 +195,12 @@ impl< mut layouter: impl Layouter, ) -> Result, D, T, RATE>, Error> { - let mode = Squeezing(poseidon_sponge( + let mode = poseidon_sponge( &self.chip, layouter.namespace(|| "PoseidonSponge"), &mut self.state, - &self.mode.0, - )?); + Some(&self.mode), + )?; Ok(Sponge { chip: self.chip, @@ -228,16 +230,12 @@ impl< } // We've already squeezed out all available elements - self.mode = Squeezing(poseidon_sponge( + self.mode = poseidon_sponge( &self.chip, layouter.namespace(|| "PoseidonSponge"), &mut self.state, - &(0..RATE) - .map(|_| Some(PaddedWord::Padding(F::zero()))) - .collect::>() - .try_into() - .unwrap(), - )?); + None, + )?; } } } diff --git a/src/circuit/gadget/poseidon/pow5.rs b/src/circuit/gadget/poseidon/pow5.rs index f6c5b865..9a74a8ac 100644 --- a/src/circuit/gadget/poseidon/pow5.rs +++ b/src/circuit/gadget/poseidon/pow5.rs @@ -9,8 +9,11 @@ use halo2::{ }; use super::{PaddedWord, PoseidonInstructions, PoseidonSpongeInstructions}; -use crate::circuit::gadget::utilities::Var; -use crate::primitives::poseidon::{Domain, Mds, Spec, SpongeRate, State}; +use crate::primitives::poseidon::{Domain, Mds, Spec, State}; +use crate::{ + circuit::gadget::utilities::Var, + primitives::poseidon::{Absorbing, Squeezing}, +}; /// Configuration for a [`Pow5Chip`]. #[derive(Clone, Debug)] @@ -313,7 +316,7 @@ impl< &self, layouter: &mut impl Layouter, initial_state: &State, - input: &SpongeRate, RATE>, + input: &Absorbing, RATE>, ) -> Result, Error> { let config = self.config(); layouter.assign_region( @@ -339,7 +342,7 @@ impl< // Load the input into this region. let load_input_word = |i: usize| { - let constraint_var = match input[i].clone() { + let constraint_var = match input.0[i].clone() { Some(PaddedWord::Message(word)) => word, Some(PaddedWord::Padding(padding_value)) => region.assign_fixed( || format!("load pad_{}", i), @@ -387,13 +390,15 @@ impl< ) } - fn get_output(state: &State) -> SpongeRate { - state[..RATE] - .iter() - .map(|word| Some(word.clone())) - .collect::>() - .try_into() - .unwrap() + fn get_output(state: &State) -> Squeezing { + Squeezing( + state[..RATE] + .iter() + .map(|word| Some(word.clone())) + .collect::>() + .try_into() + .unwrap(), + ) } } diff --git a/src/primitives/poseidon.rs b/src/primitives/poseidon.rs index 87dff7c1..34eb3c8c 100644 --- a/src/primitives/poseidon.rs +++ b/src/primitives/poseidon.rs @@ -126,14 +126,16 @@ pub(crate) fn permute, const T: usize, const RA fn poseidon_sponge, const T: usize, const RATE: usize>( state: &mut State, - input: &SpongeRate, + input: Option<&Absorbing>, mds_matrix: &Mds, round_constants: &[[F; T]], -) -> SpongeRate { - // `Iterator::zip` short-circuits when one iterator completes, so this will only - // mutate the rate portion of the state. - for (word, value) in state.iter_mut().zip(input.iter()) { - *word += value.expect("poseidon_sponge is called with a padded input"); +) -> Squeezing { + if let Some(Absorbing(input)) = input { + // `Iterator::zip` short-circuits when one iterator completes, so this will only + // mutate the rate portion of the state. + for (word, value) in state.iter_mut().zip(input.iter()) { + *word += value.expect("poseidon_sponge is called with a padded input"); + } } permute::(state, mds_matrix, round_constants); @@ -142,7 +144,7 @@ fn poseidon_sponge, const T: usize, const RATE: for (word, value) in output.iter_mut().zip(state.iter()) { *word = Some(*value); } - output + Squeezing(output) } /// The state of the [`Sponge`]. @@ -219,7 +221,7 @@ impl, const T: usize, const RATE: usize> // We've already absorbed as many elements as we can let _ = poseidon_sponge::( &mut self.state, - &self.mode.0, + Some(&self.mode), &self.mds_matrix, &self.round_constants, ); @@ -228,12 +230,12 @@ impl, const T: usize, const RATE: usize> /// Transitions the sponge into its squeezing state. pub(crate) fn finish_absorbing(mut self) -> Sponge, T, RATE> { - let mode = Squeezing(poseidon_sponge::( + let mode = poseidon_sponge::( &mut self.state, - &self.mode.0, + Some(&self.mode), &self.mds_matrix, &self.round_constants, - )); + ); Sponge { mode, @@ -258,12 +260,12 @@ impl, const T: usize, const RATE: usize> } // We've already squeezed out all available elements - self.mode = Squeezing(poseidon_sponge::( + self.mode = poseidon_sponge::( &mut self.state, - &self.mode.0, + None, &self.mds_matrix, &self.round_constants, - )); + ); } } }