From b1835798867ae64f52fc8798d4881d80945acdc2 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Fri, 9 Apr 2021 12:53:33 -0700
Subject: [PATCH] Finish up recursive Merkle proofs

---
 src/hash.rs             | 46 ++++++++++++++++++++++++++++++++++++++---
 src/merkle_proofs.rs    |  2 --
 src/plonk_challenger.rs |  1 +
 src/proof.rs            |  7 ++++++-
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/hash.rs b/src/hash.rs
index da7f1765..bb59957d 100644
--- a/src/hash.rs
+++ b/src/hash.rs
@@ -46,7 +46,48 @@ impl<F: Field> CircuitBuilder<F> {
     }
 
     pub fn hash_n_to_hash(&mut self, inputs: Vec<Target>, pad: bool) -> HashTarget {
-        todo!()
+        HashTarget::from_vec(self.hash_n_to_m(inputs, 4, pad))
+    }
+
+    pub fn hash_n_to_m(
+        &mut self,
+        mut inputs: Vec<Target>,
+        num_outputs: usize,
+        pad: bool,
+    ) -> Vec<Target> {
+        let zero = self.zero();
+        let one = self.one();
+
+        if pad {
+            inputs.push(zero);
+            while (inputs.len() + 1) % SPONGE_WIDTH != 0 {
+                inputs.push(one);
+            }
+            inputs.push(zero);
+        }
+
+        let mut state = [zero; SPONGE_WIDTH];
+
+        // Absorb all input chunks.
+        for input_chunk in inputs.chunks(SPONGE_WIDTH - 1) {
+            for i in 0..input_chunk.len() {
+                // TODO: These adds are wasteful. Maybe GMiMCGate should have separates wires to be added in.
+                state[i] = self.add(state[i], input_chunk[i]);
+            }
+            state = self.permute(state);
+        }
+
+        // Squeeze until we have the desired number of outputs.
+        let mut outputs = Vec::new();
+        loop {
+            for i in 0..(SPONGE_WIDTH - 1) {
+                outputs.push(state[i]);
+                if outputs.len() == num_outputs {
+                    return outputs;
+                }
+            }
+            state = self.permute(state);
+        }
     }
 }
 
@@ -98,8 +139,7 @@ pub fn hash_n_to_m<F: Field>(mut inputs: Vec<F>, num_outputs: usize, pad: bool)
 }
 
 pub fn hash_n_to_hash<F: Field>(inputs: Vec<F>, pad: bool) -> Hash<F> {
-    let elements = hash_n_to_m(inputs, 4, pad).try_into().unwrap();
-    Hash { elements }
+    Hash::from_vec(hash_n_to_m(inputs, 4, pad))
 }
 
 pub fn hash_n_to_1<F: Field>(inputs: Vec<F>, pad: bool) -> F {
diff --git a/src/merkle_proofs.rs b/src/merkle_proofs.rs
index 7e1b5ecf..5f735e35 100644
--- a/src/merkle_proofs.rs
+++ b/src/merkle_proofs.rs
@@ -1,5 +1,3 @@
-use std::convert::TryInto;
-
 use crate::circuit_builder::CircuitBuilder;
 use crate::field::field::Field;
 use crate::gates::gmimc::GMiMCGate;
diff --git a/src/plonk_challenger.rs b/src/plonk_challenger.rs
index 09276105..4a5a99f1 100644
--- a/src/plonk_challenger.rs
+++ b/src/plonk_challenger.rs
@@ -179,6 +179,7 @@ impl RecursiveChallenger {
         for input_chunk in self.input_buffer.chunks(SPONGE_RATE) {
             // Add the inputs to our sponge state.
             for (i, &input) in input_chunk.iter().enumerate() {
+                // TODO: These adds are wasteful. Maybe GMiMCGate should have separates wires to be added in.
                 self.sponge_state[i] = builder.add(self.sponge_state[i], input);
             }
 
diff --git a/src/proof.rs b/src/proof.rs
index 88513ce8..a77d1683 100644
--- a/src/proof.rs
+++ b/src/proof.rs
@@ -10,6 +10,11 @@ pub struct Hash<F: Field> {
 }
 
 impl<F: Field> Hash<F> {
+    pub(crate) fn from_vec(elements: Vec<F>) -> Self {
+        debug_assert!(elements.len() == 4);
+        Self { elements: elements.try_into().unwrap() }
+    }
+
     pub(crate) fn from_partial(mut elements: Vec<F>) -> Self {
         debug_assert!(elements.len() <= 4);
         while elements.len() < 4 {
@@ -27,7 +32,7 @@ pub struct HashTarget {
 impl HashTarget {
     pub(crate) fn from_vec(elements: Vec<Target>) -> Self {
         debug_assert!(elements.len() == 4);
-        HashTarget { elements: elements.try_into().unwrap() }
+        Self { elements: elements.try_into().unwrap() }
     }
 
     pub(crate) fn from_partial(mut elements: Vec<Target>, zero: Target) -> Self {