poseidon: Replace the `Duplex` struct with a `Sponge` struct

The `Sponge` struct's API correctly enforces the properties of a sponge:
it can absorb an arbitrary number of elements, and then squeeze an
arbitrary number of elements, but cannot absorb after it has squeezed.

Co-authored-by: ying tong <yingtong@z.cash>

src/circuit/gadget/poseidon.rs

@@ -3,6 +3,7 @@
 use std::array;
 use std::convert::TryInto;
 use std::fmt;
+use std::marker::PhantomData;
 
 use halo2::{
     arithmetic::FieldExt,
@@ -13,7 +14,9 @@ use halo2::{
 mod pow5;
 pub use pow5::{Pow5Chip, Pow5Config, StateWord};
 
-use crate::primitives::poseidon::{ConstantLength, Domain, Spec, Sponge, SpongeRate, State};
+use crate::primitives::poseidon::{
+    Absorbing, ConstantLength, Domain, Spec, SpongeMode, SpongeRate, Squeezing, State,
+};
 
 /// The set of circuit instructions required to use the Poseidon permutation.
 pub trait PoseidonInstructions<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>:
@@ -30,10 +33,10 @@ pub trait PoseidonInstructions<F: FieldExt, S: Spec<F, T, RATE>, const T: usize,
     ) -> Result<State<Self::Word, T>, Error>;
 }
 
-/// The set of circuit instructions required to use the [`Duplex`] and [`Hash`] gadgets.
+/// The set of circuit instructions required to use the [`Sponge`] and [`Hash`] gadgets.
 ///
 /// [`Hash`]: self::Hash
-pub trait PoseidonDuplexInstructions<
+pub trait PoseidonSpongeInstructions<
     F: FieldExt,
     S: Spec<F, T, RATE>,
     const T: usize,
@@ -91,9 +94,9 @@ impl<
     }
 }
 
-fn poseidon_duplex<
+fn poseidon_sponge<
     F: FieldExt,
-    PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
+    PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
     S: Spec<F, T, RATE>,
     D: Domain<F, T, RATE>,
     const T: usize,
@@ -110,30 +113,32 @@ fn poseidon_duplex<
     Ok(PoseidonChip::get_output(state))
 }
 
-/// A Poseidon duplex sponge.
+/// A Poseidon sponge.
 #[derive(Debug)]
-pub struct Duplex<
+pub struct Sponge<
     F: FieldExt,
-    PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
+    PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
     S: Spec<F, T, RATE>,
+    M: SpongeMode,
     D: Domain<F, T, RATE>,
     const T: usize,
     const RATE: usize,
 > {
     chip: PoseidonChip,
-    sponge: Sponge<PoseidonChip::Word, RATE>,
+    mode: M,
     state: State<PoseidonChip::Word, T>,
     domain: D,
+    _marker: PhantomData<M>,
 }
 
 impl<
         F: FieldExt,
-        PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
+        PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
         S: Spec<F, T, RATE>,
         D: Domain<F, T, RATE>,
         const T: usize,
         const RATE: usize,
-    > Duplex<F, PoseidonChip, S, D, T, RATE>
+    > Sponge<F, PoseidonChip, S, Absorbing<PoseidonChip::Word, RATE>, D, T, RATE>
 {
     /// Constructs a new duplex sponge for the given Poseidon specification.
     pub fn new(
@@ -142,9 +147,9 @@ impl<
         domain: D,
     ) -> Result<Self, Error> {
         chip.initial_state(&mut layouter, &domain)
-            .map(|state| Duplex {
+            .map(|state| Sponge {
                 chip,
-                sponge: Sponge::Absorbing(
+                mode: Absorbing(
                     (0..RATE)
                         .map(|_| None)
                         .collect::<Vec<_>>()
@@ -153,6 +158,7 @@ impl<
                 ),
                 state,
                 domain,
+                _marker: PhantomData::default(),
             })
     }
 
@@ -162,84 +168,97 @@ impl<
         mut layouter: impl Layouter<F>,
         value: AssignedCell<F, F>,
     ) -> Result<(), Error> {
-        match self.sponge {
+        for entry in self.mode.0.iter_mut() {
-            Sponge::Absorbing(ref mut input) => {
+            if entry.is_none() {
-                for entry in input.iter_mut() {
+                *entry = Some(value.into());
-                    if entry.is_none() {
+                return Ok(());
-                        *entry = Some(value.into());
-                        return Ok(());
-                    }
-                }
-
-                // We've already absorbed as many elements as we can
-                let _ = poseidon_duplex(
-                    &self.chip,
-                    layouter.namespace(|| "PoseidonDuplex"),
-                    &self.domain,
-                    &mut self.state,
-                    input,
-                )?;
-                self.sponge = Sponge::absorb(value.into());
-            }
-            Sponge::Squeezing(_) => {
-                // Drop the remaining output elements
-                self.sponge = Sponge::absorb(value.into());
             }
         }
 
+        // We've already absorbed as many elements as we can
+        let _ = poseidon_sponge(
+            &self.chip,
+            layouter.namespace(|| "PoseidonSponge"),
+            &self.domain,
+            &mut self.state,
+            &self.mode.0,
+        )?;
+        self.mode = Absorbing::init_with(value.into());
+
         Ok(())
     }
 
+    /// Transitions the sponge into its squeezing state.
+    #[allow(clippy::type_complexity)]
+    pub fn finish_absorbing(
+        mut self,
+        mut layouter: impl Layouter<F>,
+    ) -> Result<Sponge<F, PoseidonChip, S, Squeezing<PoseidonChip::Word, RATE>, D, T, RATE>, Error>
+    {
+        let mode = Squeezing(poseidon_sponge(
+            &self.chip,
+            layouter.namespace(|| "PoseidonSponge"),
+            &self.domain,
+            &mut self.state,
+            &self.mode.0,
+        )?);
+
+        Ok(Sponge {
+            chip: self.chip,
+            mode,
+            state: self.state,
+            domain: self.domain,
+            _marker: PhantomData::default(),
+        })
+    }
+}
+
+impl<
+        F: FieldExt,
+        PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
+        S: Spec<F, T, RATE>,
+        D: Domain<F, T, RATE>,
+        const T: usize,
+        const RATE: usize,
+    > Sponge<F, PoseidonChip, S, Squeezing<PoseidonChip::Word, RATE>, D, T, RATE>
+{
     /// Squeezes an element from the sponge.
     pub fn squeeze(&mut self, mut layouter: impl Layouter<F>) -> Result<AssignedCell<F, F>, Error> {
         loop {
-            match self.sponge {
+            for entry in self.mode.0.iter_mut() {
-                Sponge::Absorbing(ref input) => {
+                if let Some(inner) = entry.take() {
-                    self.sponge = Sponge::Squeezing(poseidon_duplex(
+                    return Ok(inner.into());
-                        &self.chip,
-                        layouter.namespace(|| "PoseidonDuplex"),
-                        &self.domain,
-                        &mut self.state,
-                        input,
-                    )?);
-                }
-                Sponge::Squeezing(ref mut output) => {
-                    for entry in output.iter_mut() {
-                        if let Some(inner) = entry.take() {
-                            return Ok(inner.into());
-                        }
-                    }
-
-                    // We've already squeezed out all available elements
-                    self.sponge = Sponge::Absorbing(
-                        (0..RATE)
-                            .map(|_| None)
-                            .collect::<Vec<_>>()
-                            .try_into()
-                            .unwrap(),
-                    );
                 }
             }
+
+            // We've already squeezed out all available elements
+            self.mode = Squeezing(poseidon_sponge(
+                &self.chip,
+                layouter.namespace(|| "PoseidonSponge"),
+                &self.domain,
+                &mut self.state,
+                &self.mode.0,
+            )?);
         }
     }
 }
 
-/// A Poseidon hash function, built around a duplex sponge.
+/// A Poseidon hash function, built around a sponge.
 #[derive(Debug)]
 pub struct Hash<
     F: FieldExt,
-    PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
+    PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
     S: Spec<F, T, RATE>,
     D: Domain<F, T, RATE>,
     const T: usize,
     const RATE: usize,
 > {
-    duplex: Duplex<F, PoseidonChip, S, D, T, RATE>,
+    sponge: Sponge<F, PoseidonChip, S, Absorbing<PoseidonChip::Word, RATE>, D, T, RATE>,
 }
 
 impl<
         F: FieldExt,
-        PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
+        PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
         S: Spec<F, T, RATE>,
         D: Domain<F, T, RATE>,
         const T: usize,
@@ -248,13 +267,13 @@ impl<
 {
     /// Initializes a new hasher.
     pub fn init(chip: PoseidonChip, layouter: impl Layouter<F>, domain: D) -> Result<Self, Error> {
-        Duplex::new(chip, layouter, domain).map(|duplex| Hash { duplex })
+        Sponge::new(chip, layouter, domain).map(|sponge| Hash { sponge })
     }
 }
 
 impl<
         F: FieldExt,
-        PoseidonChip: PoseidonDuplexInstructions<F, S, T, RATE>,
+        PoseidonChip: PoseidonSpongeInstructions<F, S, T, RATE>,
         S: Spec<F, T, RATE>,
         const T: usize,
         const RATE: usize,
@@ -268,9 +287,11 @@ impl<
         message: [AssignedCell<F, F>; L],
     ) -> Result<AssignedCell<F, F>, Error> {
         for (i, value) in array::IntoIter::new(message).enumerate() {
-            self.duplex
+            self.sponge
                 .absorb(layouter.namespace(|| format!("absorb_{}", i)), value)?;
         }
-        self.duplex.squeeze(layouter.namespace(|| "squeeze"))
+        self.sponge
+            .finish_absorbing(layouter.namespace(|| "finish absorbing"))?
+            .squeeze(layouter.namespace(|| "squeeze"))
     }
 }

src/circuit/gadget/poseidon/pow5.rs

@@ -8,7 +8,7 @@ use halo2::{
     poly::Rotation,
 };
 
-use super::{PoseidonDuplexInstructions, PoseidonInstructions};
+use super::{PoseidonInstructions, PoseidonSpongeInstructions};
 use crate::circuit::gadget::utilities::Var;
 use crate::primitives::poseidon::{Domain, Mds, Spec, SpongeRate, State};
 
@@ -269,7 +269,7 @@ impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize
 }
 
 impl<F: FieldExt, S: Spec<F, WIDTH, RATE>, const WIDTH: usize, const RATE: usize>
-    PoseidonDuplexInstructions<F, S, WIDTH, RATE> for Pow5Chip<F, WIDTH, RATE>
+    PoseidonSpongeInstructions<F, S, WIDTH, RATE> for Pow5Chip<F, WIDTH, RATE>
 {
     fn initial_state(
         &self,

src/primitives/poseidon.rs

@@ -24,7 +24,7 @@ use grain::SboxType;
 /// The type used to hold permutation state.
 pub(crate) type State<F, const T: usize> = [F; T];
 
-/// The type used to hold duplex sponge state.
+/// The type used to hold sponge rate.
 pub(crate) type SpongeRate<F, const RATE: usize> = [Option<F>; RATE];
 
 /// The type used to hold the MDS matrix and its inverse.
@@ -124,7 +124,7 @@ pub(crate) fn permute<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RA
         });
 }
 
-fn poseidon_duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>(
+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>,
     pad_and_add: &dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>),
@@ -142,48 +142,65 @@ fn poseidon_duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE:
     output
 }
 
-#[derive(Debug)]
+/// The state of the [`Sponge`].
-pub(crate) enum Sponge<F, const RATE: usize> {
+// TODO: Seal this trait?
-    Absorbing(SpongeRate<F, RATE>),
+pub trait SpongeMode {}
-    Squeezing(SpongeRate<F, RATE>),
-}
 
-impl<F: fmt::Debug, const RATE: usize> Sponge<F, RATE> {
+/// The absorbing state of the [`Sponge`].
-    pub(crate) fn absorb(val: F) -> Self {
+#[derive(Debug)]
-        let mut input: [Option<F>; RATE] = (0..RATE)
+pub struct Absorbing<F, const RATE: usize>(pub(crate) SpongeRate<F, RATE>);
-            .map(|_| None)
+
-            .collect::<Vec<_>>()
+/// The squeezing state of the [`Sponge`].
-            .try_into()
+#[derive(Debug)]
-            .unwrap();
+pub struct Squeezing<F, const RATE: usize>(pub(crate) SpongeRate<F, RATE>);
-        input[0] = Some(val);
+
-        Sponge::Absorbing(input)
+impl<F, const RATE: usize> SpongeMode for Absorbing<F, RATE> {}
+impl<F, const RATE: usize> SpongeMode for Squeezing<F, RATE> {}
+
+impl<F: fmt::Debug, const RATE: usize> Absorbing<F, RATE> {
+    pub(crate) fn init_with(val: F) -> Self {
+        Self(
+            iter::once(Some(val))
+                .chain((1..RATE).map(|_| None))
+                .collect::<Vec<_>>()
+                .try_into()
+                .unwrap(),
+        )
     }
 }
 
-/// A Poseidon duplex sponge.
+/// A Poseidon sponge.
-pub(crate) struct Duplex<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> {
+pub(crate) struct Sponge<
-    sponge: Sponge<F, RATE>,
+    F: FieldExt,
+    S: Spec<F, T, RATE>,
+    M: SpongeMode,
+    const T: usize,
+    const RATE: usize,
+> {
+    mode: M,
     state: State<F, T>,
     pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>)>,
     mds_matrix: Mds<F, T>,
     round_constants: Vec<[F; T]>,
-    _marker: PhantomData<S>,
+    _marker: PhantomData<(S, M)>,
 }
 
-impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> Duplex<F, S, T, RATE> {
+impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
-    /// Constructs a new duplex sponge for the given Poseidon specification.
+    Sponge<F, S, Absorbing<F, RATE>, T, RATE>
+{
+    /// Constructs a new sponge for the given Poseidon specification.
     pub(crate) fn new(
         initial_capacity_element: F,
         pad_and_add: Box<dyn Fn(&mut State<F, T>, &SpongeRate<F, RATE>)>,
     ) -> Self {
         let (round_constants, mds_matrix, _) = S::constants();
 
-        let input = [None; RATE];
+        let mode = Absorbing([None; RATE]);
         let mut state = [F::zero(); T];
         state[RATE] = initial_capacity_element;
 
-        Duplex {
+        Sponge {
-            sponge: Sponge::Absorbing(input),
+            mode,
             state,
             pad_and_add,
             mds_matrix,
@@ -194,56 +211,65 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize> Duplex
 
     /// Absorbs an element into the sponge.
     pub(crate) fn absorb(&mut self, value: F) {
-        match self.sponge {
+        for entry in self.mode.0.iter_mut() {
-            Sponge::Absorbing(ref mut input) => {
+            if entry.is_none() {
-                for entry in input.iter_mut() {
+                *entry = Some(value);
-                    if entry.is_none() {
+                return;
-                        *entry = Some(value);
-                        return;
-                    }
-                }
-
-                // We've already absorbed as many elements as we can
-                let _ = poseidon_duplex::<F, S, T, RATE>(
-                    &mut self.state,
-                    input,
-                    &self.pad_and_add,
-                    &self.mds_matrix,
-                    &self.round_constants,
-                );
-                self.sponge = Sponge::absorb(value);
-            }
-            Sponge::Squeezing(_) => {
-                // Drop the remaining output elements
-                self.sponge = Sponge::absorb(value);
             }
         }
+
+        // We've already absorbed as many elements as we can
+        let _ = poseidon_sponge::<F, S, T, RATE>(
+            &mut self.state,
+            &self.mode.0,
+            &self.pad_and_add,
+            &self.mds_matrix,
+            &self.round_constants,
+        );
+        self.mode = Absorbing::init_with(value);
     }
 
+    /// 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>(
+            &mut self.state,
+            &self.mode.0,
+            &self.pad_and_add,
+            &self.mds_matrix,
+            &self.round_constants,
+        ));
+
+        Sponge {
+            mode,
+            state: self.state,
+            pad_and_add: self.pad_and_add,
+            mds_matrix: self.mds_matrix,
+            round_constants: self.round_constants,
+            _marker: PhantomData::default(),
+        }
+    }
+}
+
+impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize>
+    Sponge<F, S, Squeezing<F, RATE>, T, RATE>
+{
     /// Squeezes an element from the sponge.
     pub(crate) fn squeeze(&mut self) -> F {
         loop {
-            match self.sponge {
+            for entry in self.mode.0.iter_mut() {
-                Sponge::Absorbing(ref input) => {
+                if let Some(e) = entry.take() {
-                    self.sponge = Sponge::Squeezing(poseidon_duplex::<F, S, T, RATE>(
+                    return e;
-                        &mut self.state,
-                        input,
-                        &self.pad_and_add,
-                        &self.mds_matrix,
-                        &self.round_constants,
-                    ));
-                }
-                Sponge::Squeezing(ref mut output) => {
-                    for entry in output.iter_mut() {
-                        if let Some(e) = entry.take() {
-                            return e;
-                        }
-                    }
-
-                    // We've already squeezed out all available elements
-                    self.sponge = Sponge::Absorbing([None; RATE]);
                 }
             }
+
+            // We've already squeezed out all available elements
+            self.mode = Squeezing(poseidon_sponge::<F, S, T, RATE>(
+                &mut self.state,
+                &self.mode.0,
+                &self.pad_and_add,
+                &self.mds_matrix,
+                &self.round_constants,
+            ));
         }
     }
 }
@@ -301,7 +327,7 @@ impl<F: FieldExt, const T: usize, const RATE: usize, const L: usize> Domain<F, T
     }
 }
 
-/// A Poseidon hash function, built around a duplex sponge.
+/// A Poseidon hash function, built around a sponge.
 pub struct Hash<
     F: FieldExt,
     S: Spec<F, T, RATE>,
@@ -309,7 +335,7 @@ pub struct Hash<
     const T: usize,
     const RATE: usize,
 > {
-    duplex: Duplex<F, S, T, RATE>,
+    sponge: Sponge<F, S, Absorbing<F, RATE>, T, RATE>,
     domain: D,
 }
 
@@ -343,7 +369,7 @@ impl<
     /// Initializes a new hasher.
     pub fn init(domain: D) -> Self {
         Hash {
-            duplex: Duplex::new(domain.initial_capacity_element(), domain.pad_and_add()),
+            sponge: Sponge::new(domain.initial_capacity_element(), domain.pad_and_add()),
             domain,
         }
     }
@@ -355,9 +381,9 @@ impl<F: FieldExt, S: Spec<F, T, RATE>, const T: usize, const RATE: usize, const
     /// Hashes the given input.
     pub fn hash(mut self, message: [F; L]) -> F {
         for value in array::IntoIter::new(message) {
-            self.duplex.absorb(value);
+            self.sponge.absorb(value);
         }
-        self.duplex.squeeze()
+        self.sponge.finish_absorbing().squeeze()
     }
 }