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

@@ -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<F>,
         initial_state: &State<Self::Word, T>,
-        input: &SpongeRate<PaddedWord<F>, RATE>,
+        input: &Absorbing<PaddedWord<F>, RATE>,
     ) -> Result<State<Self::Word, T>, Error>;
 
     /// Extracts sponge output from the given state.
-    fn get_output(state: &State<Self::Word, T>) -> SpongeRate<Self::Word, RATE>;
+    fn get_output(state: &State<Self::Word, T>) -> Squeezing<Self::Word, RATE>;
 }
 
 /// A word over which the Poseidon permutation operates.
@@ -112,9 +112,11 @@ fn poseidon_sponge<
     chip: &PoseidonChip,
     mut layouter: impl Layouter<F>,
     state: &mut State<PoseidonChip::Word, T>,
-    input: &SpongeRate<PaddedWord<F>, RATE>,
-) -> Result<SpongeRate<PoseidonChip::Word, RATE>, Error> {
-    *state = chip.add_input(&mut layouter, state, input)?;
+    input: Option<&Absorbing<PaddedWord<F>, RATE>>,
+) -> Result<Squeezing<PoseidonChip::Word, RATE>, 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<F>,
     ) -> Result<Sponge<F, PoseidonChip, S, Squeezing<PoseidonChip::Word, RATE>, 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::<Vec<_>>()
-                    .try_into()
-                    .unwrap(),
-            )?);
+                None,
+            )?;
         }
     }
 }

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<F>,
         initial_state: &State<Self::Word, WIDTH>,
-        input: &SpongeRate<PaddedWord<F>, RATE>,
+        input: &Absorbing<PaddedWord<F>, RATE>,
     ) -> Result<State<Self::Word, WIDTH>, 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<Self::Word, WIDTH>) -> SpongeRate<Self::Word, RATE> {
-        state[..RATE]
-            .iter()
-            .map(|word| Some(word.clone()))
-            .collect::<Vec<_>>()
-            .try_into()
-            .unwrap()
+    fn get_output(state: &State<Self::Word, WIDTH>) -> Squeezing<Self::Word, RATE> {
+        Squeezing(
+            state[..RATE]
+                .iter()
+                .map(|word| Some(word.clone()))
+                .collect::<Vec<_>>()
+                .try_into()
+                .unwrap(),
+        )
     }
 }
 

src/primitives/poseidon.rs

@@ -126,14 +126,16 @@ pub(crate) fn permute<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RA
 
 fn poseidon_sponge<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
     state: &mut State<F, T>,
-    input: &SpongeRate<F, RATE>,
+    input: Option<&Absorbing<F, RATE>>,
     mds_matrix: &Mds<F, T>,
     round_constants: &[[F; T]],
-) -> SpongeRate<F, RATE> {
-    // `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<F, RATE> {
+    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::<F, S, T, RATE>(state, mds_matrix, round_constants);
@@ -142,7 +144,7 @@ fn poseidon_sponge<F: FieldExt, S: Spec<F, T, RATE>, 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<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
         // We've already absorbed as many elements as we can
         let _ = poseidon_sponge::<F, S, T, RATE>(
             &mut self.state,
-            &self.mode.0,
+            Some(&self.mode),
             &self.mds_matrix,
             &self.round_constants,
         );
@@ -228,12 +230,12 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
 
     /// Transitions the sponge into its squeezing state.
     pub(crate) fn finish_absorbing(mut self) -> Sponge<F, S, Squeezing<F, RATE>, T, RATE> {
-        let mode = Squeezing(poseidon_sponge::<F, S, T, RATE>(
+        let mode = poseidon_sponge::<F, S, T, RATE>(
             &mut self.state,
-            &self.mode.0,
+            Some(&self.mode),
             &self.mds_matrix,
             &self.round_constants,
-        ));
+        );
 
         Sponge {
             mode,
@@ -258,12 +260,12 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
             }
 
             // We've already squeezed out all available elements
-            self.mode = Squeezing(poseidon_sponge::<F, S, T, RATE>(
+            self.mode = poseidon_sponge::<F, S, T, RATE>(
                 &mut self.state,
-                &self.mode.0,
+                None,
                 &self.mds_matrix,
                 &self.round_constants,
-            ));
+            );
         }
     }
 }