From 266705166f6fa941d971f2714fb67a42dce4eb78 Mon Sep 17 00:00:00 2001
From: Jack Grigg <thestr4d@gmail.com>
Date: Thu, 11 Mar 2021 18:35:45 +1300
Subject: [PATCH] Poseidon duplex sponge and hash function

---
 src/primitives/poseidon.rs | 176 +++++++++++++++++++++++++++++++++++++
 1 file changed, 176 insertions(+)

diff --git a/src/primitives/poseidon.rs b/src/primitives/poseidon.rs
index 7aa411ab..89b9554c 100644
--- a/src/primitives/poseidon.rs
+++ b/src/primitives/poseidon.rs
@@ -1,3 +1,4 @@
+use std::iter;
 use std::marker::PhantomData;
 
 use halo2::arithmetic::FieldExt;
@@ -92,3 +93,178 @@ impl<F: FieldExt> Spec<F> for Generic<F> {
         (round_constants, mds, mds_inv)
     }
 }
+
+/// Runs the Poseidon permutation on the given state.
+fn permute<F: FieldExt, S: Spec<F>>(
+    state: &mut [F],
+    spec: &S,
+    mds: &[Vec<F>],
+    round_constants: &[Vec<F>],
+) {
+    // TODO: Remove this when we can use const generics.
+    assert!(state.len() == spec.arity());
+
+    let r_f = spec.full_rounds() / 2;
+    let r_p = spec.partial_rounds();
+
+    let apply_mds = |state: &mut [F]| {
+        let new_state: Vec<_> = mds
+            .iter()
+            .map(|mds_row| {
+                mds_row
+                    .iter()
+                    .zip(state.iter())
+                    .fold(F::zero(), |acc, (mds, word)| acc + *mds * *word)
+            })
+            .collect();
+        for (word, new_word) in state.iter_mut().zip(new_state.into_iter()) {
+            *word = new_word;
+        }
+    };
+
+    let full_round = |state: &mut [F], rcs: &[F]| {
+        for (word, rc) in state.iter_mut().zip(rcs.iter()) {
+            *word = spec.sbox(*word + rc);
+        }
+        apply_mds(state);
+    };
+
+    let part_round = |state: &mut [F], rcs: &[F]| {
+        for (word, rc) in state.iter_mut().zip(rcs.iter()) {
+            *word += rc;
+        }
+        state[0] = spec.sbox(state[0]);
+        apply_mds(state);
+    };
+
+    iter::empty()
+        .chain(iter::repeat(&full_round as &dyn Fn(&mut [F], &[F])).take(r_f))
+        .chain(iter::repeat(&part_round as &dyn Fn(&mut [F], &[F])).take(r_p))
+        .chain(iter::repeat(&full_round as &dyn Fn(&mut [F], &[F])).take(r_f))
+        .zip(round_constants.iter())
+        .fold(state, |state, (round, rcs)| {
+            round(state, rcs);
+            state
+        });
+}
+
+fn pad_and_add<F: FieldExt>(state: &mut [F], input: &[F]) {
+    let padding = state.len() - input.len();
+    // TODO: Decide on a padding strategy (currently padding with all-ones)
+    for (word, val) in state
+        .iter_mut()
+        .zip(input.iter().chain(iter::repeat(&F::one()).take(padding)))
+    {
+        *word += val;
+    }
+}
+
+enum SpongeState<F: FieldExt> {
+    Absorbing(Vec<F>),
+    Squeezing(Vec<F>),
+}
+
+/// A Poseidon duplex sponge.
+pub struct Duplex<F: FieldExt, S: Spec<F>> {
+    spec: S,
+    sponge: Option<SpongeState<F>>,
+    state: Vec<F>,
+    rate: usize,
+    mds_matrix: Vec<Vec<F>>,
+    round_constants: Vec<Vec<F>>,
+}
+
+impl<F: FieldExt, S: Spec<F>> Duplex<F, S> {
+    /// Constructs a new duplex sponge with the given rate.
+    pub fn new(spec: S, rate: usize) -> Self {
+        assert!(rate < spec.arity());
+
+        let state = vec![F::zero(); spec.arity()];
+        let (round_constants, mds_matrix, _) = spec.constants();
+
+        Duplex {
+            spec,
+            sponge: Some(SpongeState::Absorbing(vec![])),
+            state,
+            rate,
+            mds_matrix,
+            round_constants,
+        }
+    }
+
+    fn process(&mut self, input: &[F]) -> Vec<F> {
+        pad_and_add(&mut self.state[..self.rate], input);
+
+        permute(
+            &mut self.state,
+            &self.spec,
+            &self.mds_matrix,
+            &self.round_constants,
+        );
+
+        self.state[..self.rate].to_vec()
+    }
+
+    /// Absorbs an element into the sponge.
+    pub fn absorb(&mut self, value: F) {
+        match self.sponge.take().unwrap() {
+            SpongeState::Absorbing(mut input) => {
+                if input.len() < self.rate {
+                    input.push(value);
+                    self.sponge = Some(SpongeState::Absorbing(input));
+                    return;
+                }
+
+                // We've already absorbed as many elements as we can
+                let _ = self.process(&input);
+                self.sponge = Some(SpongeState::Absorbing(vec![value]));
+            }
+            SpongeState::Squeezing(_) => {
+                // Drop the remaining output elements
+                self.sponge = Some(SpongeState::Absorbing(vec![value]));
+            }
+        }
+    }
+
+    /// Squeezes an element from the sponge.
+    pub fn squeeze(&mut self) -> F {
+        loop {
+            match self.sponge.take().unwrap() {
+                SpongeState::Absorbing(input) => {
+                    self.sponge = Some(SpongeState::Squeezing(self.process(&input)));
+                }
+                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![]));
+                }
+            }
+        }
+    }
+}
+
+/// A Poseidon hash function, built around a duplex sponge.
+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 {
+        Hash(Duplex::new(spec, rate))
+    }
+
+    /// Updates the hasher with the given value.
+    pub fn update(&mut self, value: F) {
+        self.0.absorb(value);
+    }
+
+    /// Finalizes the hasher, returning its output.
+    pub fn finalize(mut self) -> F {
+        // TODO: Check which state element other implementations use.
+        self.0.squeeze()
+    }
+}