From b63d83aacff58d914da46f84c84ff8c3f53844b3 Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Thu, 16 Sep 2021 21:18:08 +0200
Subject: [PATCH] Add Poseidon gadget

---
 src/gadgets/hash.rs             | 44 ++++++++++++++++++++++++++++++---
 src/hash/hashing.rs             |  2 +-
 src/plonk/circuit_data.rs       |  2 +-
 src/plonk/recursive_verifier.rs |  2 +-
 4 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/src/gadgets/hash.rs b/src/gadgets/hash.rs
index 01519912..82de1095 100644
--- a/src/gadgets/hash.rs
+++ b/src/gadgets/hash.rs
@@ -3,8 +3,10 @@ use std::convert::TryInto;
 use crate::field::extension_field::Extendable;
 use crate::field::field_types::RichField;
 use crate::gates::gmimc::GMiMCGate;
+use crate::gates::poseidon::PoseidonGate;
 use crate::hash::gmimc::GMiMC;
 use crate::hash::hashing::{HashFamily, HASH_FAMILY};
+use crate::hash::poseidon::Poseidon;
 use crate::iop::target::{BoolTarget, Target};
 use crate::iop::wire::Wire;
 use crate::plonk::circuit_builder::CircuitBuilder;
@@ -13,7 +15,8 @@ use crate::plonk::circuit_builder::CircuitBuilder;
 impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
     pub fn permute<const W: usize>(&mut self, inputs: [Target; W]) -> [Target; W]
     where
-        F: GMiMC<W>,
+        F: GMiMC<W> + Poseidon<W>,
+        [(); W - 1]: ,
     {
         // We don't want to swap any inputs, so set that wire to 0.
         let _false = self._false();
@@ -28,11 +31,12 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
         swap: BoolTarget,
     ) -> [Target; W]
     where
-        F: GMiMC<W>,
+        F: GMiMC<W> + Poseidon<W>,
+        [(); W - 1]: ,
     {
         match HASH_FAMILY {
             HashFamily::GMiMC => self.gmimc_permute_swapped(inputs, swap),
-            HashFamily::Poseidon => todo!(),
+            HashFamily::Poseidon => self.poseidon_permute_swapped(inputs, swap),
         }
     }
 
@@ -79,4 +83,38 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilder<F, D> {
             .try_into()
             .unwrap()
     }
+
+    /// Conditionally swap two chunks of the inputs (useful in verifying Merkle proofs), then apply
+    /// the Poseidon permutation.
+    pub(crate) fn poseidon_permute_swapped<const W: usize>(
+        &mut self,
+        inputs: [Target; W],
+        swap: BoolTarget,
+    ) -> [Target; W]
+    where
+        F: Poseidon<W>,
+        [(); W - 1]: ,
+    {
+        let gate_type = PoseidonGate::<F, D, W>::new();
+        let gate = self.add_gate(gate_type, vec![]);
+
+        // We don't want to swap any inputs, so set that wire to 0.
+        let swap_wire = PoseidonGate::<F, D, W>::WIRE_SWAP;
+        let swap_wire = Target::wire(gate, swap_wire);
+        self.connect(swap.target, swap_wire);
+
+        // Route input wires.
+        for i in 0..W {
+            let in_wire = PoseidonGate::<F, D, W>::wire_input(i);
+            let in_wire = Target::wire(gate, in_wire);
+            self.connect(inputs[i], in_wire);
+        }
+
+        // Collect output wires.
+        (0..W)
+            .map(|i| Target::wire(gate, PoseidonGate::<F, D, W>::wire_output(i)))
+            .collect::<Vec<_>>()
+            .try_into()
+            .unwrap()
+    }
 }
diff --git a/src/hash/hashing.rs b/src/hash/hashing.rs
index 9b4791d3..ae14e058 100644
--- a/src/hash/hashing.rs
+++ b/src/hash/hashing.rs
@@ -10,7 +10,7 @@ pub(crate) const SPONGE_RATE: usize = 8;
 pub(crate) const SPONGE_CAPACITY: usize = 4;
 pub(crate) const SPONGE_WIDTH: usize = SPONGE_RATE + SPONGE_CAPACITY;
 
-pub(crate) const HASH_FAMILY: HashFamily = HashFamily::GMiMC;
+pub(crate) const HASH_FAMILY: HashFamily = HashFamily::Poseidon;
 
 pub(crate) enum HashFamily {
     GMiMC,
diff --git a/src/plonk/circuit_data.rs b/src/plonk/circuit_data.rs
index 21dfc28b..540e3e84 100644
--- a/src/plonk/circuit_data.rs
+++ b/src/plonk/circuit_data.rs
@@ -61,7 +61,7 @@ impl CircuitConfig {
     #[cfg(test)]
     pub(crate) fn large_config() -> Self {
         Self {
-            num_wires: 126,
+            num_wires: 143,
             num_routed_wires: 64,
             security_bits: 128,
             rate_bits: 3,
diff --git a/src/plonk/recursive_verifier.rs b/src/plonk/recursive_verifier.rs
index 3bd50bec..6b0bd8c3 100644
--- a/src/plonk/recursive_verifier.rs
+++ b/src/plonk/recursive_verifier.rs
@@ -416,7 +416,7 @@ mod tests {
         type F = CrandallField;
         const D: usize = 4;
         let config = CircuitConfig {
-            num_wires: 126,
+            num_wires: 143,
             num_routed_wires: 64,
             security_bits: 128,
             rate_bits: 3,