Add poseidon::Spec::Rate associated type

This removes the need for specifying the rate at runtime, and removes
the remaining heap allocations from Duplex::absorb and Duplex::squeeze.

src/primitives/poseidon.rs

@@ -18,6 +18,13 @@ pub trait Spec<F: FieldExt> {
     /// This must be an array of length [`Spec::arity`], that defaults to all-zeroes.
     type State: Default + AsRef<[F]> + AsMut<[F]>;
 
+    /// The type used to hold duplex sponge state.
+    ///
+    /// This must be an array of length equal to the rate of the duplex sponge (allowing
+    /// for a capacity consistent with this specification's security level), that defaults
+    /// to `[None; RATE]`.
+    type Rate: Default + AsRef<[Option<F>]> + AsMut<[Option<F>]>;
+
     /// The arity of this specification.
     fn arity() -> usize;
 
@@ -104,6 +111,7 @@ impl<F: FieldExt> P256Pow5T3<F> {
 
 impl<F: FieldExt> Spec<F> for P256Pow5T3<F> {
     type State = [F; 3];
+    type Rate = [Option<F>; 2];
 
     fn arity() -> usize {
         3
@@ -174,76 +182,88 @@ fn permute<F: FieldExt, S: Spec<F>>(
         });
 }
 
-fn pad_and_add<F: FieldExt>(state: &mut [F], input: &[F]) {
+fn poseidon_duplex<F: FieldExt, S: Spec<F>>(
-    let padding = state.len() - input.len();
+    state: &mut S::State,
-    // TODO: Decide on a padding strategy (currently padding with all-ones)
+    input: &S::Rate,
-    for (word, val) in state
+    mds_matrix: &[S::State],
-        .iter_mut()
+    round_constants: &[S::State],
-        .zip(input.iter().chain(iter::repeat(&F::one()).take(padding)))
+) -> S::Rate {
-    {
+    // `Iterator::zip` short-circuits when one iterator completes, so this will only
-        *word += val;
+    // mutate the rate portion of the state.
+    for (word, value) in state.as_mut().iter_mut().zip(input.as_ref().iter()) {
+        // TODO: Decide on a padding strategy, if we ever need to use Poseidon with
+        // incomplete state input.
+        *word += value.unwrap();
     }
+
+    permute::<F, S>(state, mds_matrix, round_constants);
+
+    let mut output = S::Rate::default();
+    for (word, value) in output.as_mut().iter_mut().zip(state.as_ref().iter()) {
+        *word = Some(*value);
+    }
+    output
 }
 
-enum SpongeState<F: FieldExt> {
+enum SpongeState<F: FieldExt, S: Spec<F>> {
-    Absorbing(Vec<F>),
+    Absorbing(S::Rate),
-    Squeezing(Vec<F>),
+    Squeezing(S::Rate),
+}
+
+impl<F: FieldExt, S: Spec<F>> SpongeState<F, S> {
+    fn absorb(val: F) -> Self {
+        let mut input = S::Rate::default();
+        input.as_mut()[0] = Some(val);
+        SpongeState::Absorbing(input)
+    }
 }
 
 /// A Poseidon duplex sponge.
 pub struct Duplex<F: FieldExt, S: Spec<F>> {
-    sponge: Option<SpongeState<F>>,
+    sponge: SpongeState<F, S>,
     state: S::State,
-    rate: usize,
     mds_matrix: Vec<S::State>,
     round_constants: Vec<S::State>,
     _marker: PhantomData<S>,
 }
 
 impl<F: FieldExt, S: Spec<F>> Duplex<F, S> {
-    /// Constructs a new duplex sponge with the given rate.
+    /// Constructs a new duplex sponge for the given Poseidon specification.
-    pub fn new(spec: S, rate: usize) -> Self {
+    pub fn new(spec: S) -> Self {
-        // The sponge capacity must be at least 1.
-        // TODO: Construct the capacity from the specification's security level.
-        assert!(rate < S::arity());
-
         let (round_constants, mds_matrix, _) = spec.constants();
 
         Duplex {
-            sponge: Some(SpongeState::Absorbing(vec![])),
+            sponge: SpongeState::Absorbing(S::Rate::default()),
             state: S::State::default(),
-            rate,
             mds_matrix,
             round_constants,
             _marker: PhantomData::default(),
         }
     }
 
-    fn process(&mut self, input: &[F]) -> Vec<F> {
-        pad_and_add(&mut self.state.as_mut()[..self.rate], input);
-
-        permute::<F, S>(&mut self.state, &self.mds_matrix, &self.round_constants);
-
-        self.state.as_ref()[..self.rate].to_vec()
-    }
-
     /// Absorbs an element into the sponge.
     pub fn absorb(&mut self, value: F) {
-        match self.sponge.take().unwrap() {
+        match self.sponge {
-            SpongeState::Absorbing(mut input) => {
+            SpongeState::Absorbing(ref mut input) => {
-                if input.len() < self.rate {
+                for entry in input.as_mut().iter_mut() {
-                    input.push(value);
+                    if entry.is_none() {
-                    self.sponge = Some(SpongeState::Absorbing(input));
+                        *entry = Some(value);
                         return;
                     }
+                }
 
                 // We've already absorbed as many elements as we can
-                let _ = self.process(&input);
+                let _ = poseidon_duplex::<F, S>(
-                self.sponge = Some(SpongeState::Absorbing(vec![value]));
+                    &mut self.state,
+                    &input,
+                    &self.mds_matrix,
+                    &self.round_constants,
+                );
+                self.sponge = SpongeState::absorb(value);
             }
             SpongeState::Squeezing(_) => {
                 // Drop the remaining output elements
-                self.sponge = Some(SpongeState::Absorbing(vec![value]));
+                self.sponge = SpongeState::absorb(value);
             }
         }
     }
@@ -251,19 +271,24 @@ impl<F: FieldExt, S: Spec<F>> Duplex<F, S> {
     /// Squeezes an element from the sponge.
     pub fn squeeze(&mut self) -> F {
         loop {
-            match self.sponge.take().unwrap() {
+            match self.sponge {
-                SpongeState::Absorbing(input) => {
+                SpongeState::Absorbing(ref input) => {
-                    self.sponge = Some(SpongeState::Squeezing(self.process(&input)));
+                    self.sponge = SpongeState::Squeezing(poseidon_duplex::<F, S>(
+                        &mut self.state,
+                        &input,
+                        &self.mds_matrix,
+                        &self.round_constants,
+                    ));
+                }
+                SpongeState::Squeezing(ref mut output) => {
+                    for entry in output.as_mut().iter_mut() {
+                        if let Some(e) = entry.take() {
+                            return e;
                         }
-                SpongeState::Squeezing(mut output) => {
-                    if !output.is_empty() {
-                        let ret = output.remove(0);
-                        self.sponge = Some(SpongeState::Squeezing(output));
-                        return ret;
                     }
 
                     // We've already squeezed out all available elements
-                    self.sponge = Some(SpongeState::Absorbing(vec![]));
+                    self.sponge = SpongeState::Absorbing(S::Rate::default());
                 }
             }
         }
@@ -275,8 +300,8 @@ pub struct Hash<F: FieldExt, S: Spec<F>>(Duplex<F, S>);
 
 impl<F: FieldExt, S: Spec<F>> Hash<F, S> {
     /// Initializes a new hasher.
-    pub fn init(spec: S, rate: usize) -> Self {
+    pub fn init(spec: S) -> Self {
-        Hash(Duplex::new(spec, rate))
+        Hash(Duplex::new(spec))
     }
 
     /// Updates the hasher with the given value.