From 03c7baedeae208b21359e542b8683644b3eb1ac2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Antonio=20Mej=C3=ADas=20Gil?= <anmegi.95@gmail.com>
Date: Mon, 31 Jul 2023 09:03:51 +0200
Subject: [PATCH] Compatibility with native Poseidon sponge (#98)

* fixed compatibility issue, added explanation

* added comment, formatted

* changed corrected computation into circuit constraint, removed old code

* added first version of test

* finalised tests

* remove unnecessary builder construction from test

* removed unnecessary variable k

* added poseidon test to ci.yml
---
 .github/workflows/ci.yml     |   4 ++
 hashes/poseidon/src/lib.rs   |  17 ++++-
 hashes/poseidon/src/tests.rs | 119 +++++++++++++++++++++++++++++++++++
 3 files changed, 138 insertions(+), 2 deletions(-)
 create mode 100644 hashes/poseidon/src/tests.rs

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index cd87d1b5..0a8c2db3 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,6 +21,10 @@ jobs:
         working-directory: 'halo2-base'
         run: |
           cargo test -- --test-threads=1
+      - name: Run poseidon tests
+        working-directory: 'hashes/poseidon'
+        run: |
+          cargo test test_poseidon_compatibility
       - name: Run halo2-ecc tests MockProver
         working-directory: 'halo2-ecc'
         run: |
diff --git a/hashes/poseidon/src/lib.rs b/hashes/poseidon/src/lib.rs
index 0d4a0701..952b8288 100644
--- a/hashes/poseidon/src/lib.rs
+++ b/hashes/poseidon/src/lib.rs
@@ -9,6 +9,8 @@ use halo2_base::{
     QuantumCell::{Constant, Existing},
 };
 
+pub mod tests;
+
 struct PoseidonState<F: ScalarField, const T: usize, const RATE: usize> {
     s: [AssignedValue<F>; T],
 }
@@ -51,15 +53,26 @@ impl<F: ScalarField, const T: usize, const RATE: usize> PoseidonState<F, T, RATE
         assert!(inputs.len() < T);
         let offset = inputs.len() + 1;
 
-        self.s[0] =
-            gate.sum(ctx, inputs.iter().map(|a| Existing(*a)).chain([Constant(pre_constants[0])]));
+        // Explanation of what's going on: before each round of the poseidon permutation,
+        // two things have to be added to the state: inputs (the absorbed elements) and
+        // preconstants. Imagine the state as a list of T elements, the first of which is
+        // the capacity:  |--cap--|--el1--|--el2--|--elR--|
+        // - A preconstant is added to each of all T elements (which is different for each)
+        // - The inputs are added to all elements starting from el1 (so, not to the capacity),
+        //   to as many elements as inputs are available.
+        // - To the first element for which no input is left (if any), an extra 1 is added.
+
+        // adding preconstant to the distinguished capacity element (only one)
+        self.s[0] = gate.add(ctx, self.s[0], Constant(pre_constants[0]));
 
+        // adding pre-constants and inputs to the elements for which both are available
         for ((x, constant), input) in
             self.s.iter_mut().skip(1).zip(pre_constants.iter().skip(1)).zip(inputs.iter())
         {
             *x = gate.sum(ctx, [Existing(*x), Existing(*input), Constant(*constant)]);
         }
 
+        // adding only pre-constants when no input is left
         for (i, (x, constant)) in
             self.s.iter_mut().skip(offset).zip(pre_constants.iter().skip(offset)).enumerate()
         {
diff --git a/hashes/poseidon/src/tests.rs b/hashes/poseidon/src/tests.rs
new file mode 100644
index 00000000..fd9625ad
--- /dev/null
+++ b/hashes/poseidon/src/tests.rs
@@ -0,0 +1,119 @@
+#[cfg(test)]
+mod tests {
+    use std::{cmp::max, iter::zip};
+
+    use halo2_base::{
+        gates::{builder::GateThreadBuilder, GateChip},
+        halo2_proofs::halo2curves::bn256::Fr,
+        utils::ScalarField,
+    };
+    use poseidon::Poseidon;
+    use rand::Rng;
+
+    use crate::PoseidonChip;
+
+    // make interleaved calls to absorb and squeeze elements and
+    // check that the result is the same in-circuit and natively
+    fn poseidon_compatiblity_verification<F: ScalarField, const T: usize, const RATE: usize>(
+        // elements of F to absorb; one sublist = one absorption
+        mut absorptions: Vec<Vec<F>>,
+        // list of amounts of elements of F that should be squeezed every time
+        mut squeezings: Vec<usize>,
+        rounds_full: usize,
+        rounts_partial: usize,
+    ) {
+        let mut builder = GateThreadBuilder::prover();
+        let gate = GateChip::default();
+
+        let mut ctx = builder.main(0);
+
+        // constructing native and in-circuit Poseidon sponges
+        let mut native_sponge = Poseidon::<F, T, RATE>::new(rounds_full, rounts_partial);
+        let mut circuit_sponge =
+            PoseidonChip::<F, T, RATE>::new(&mut ctx, rounds_full, rounts_partial)
+                .expect("Failed to construct Poseidon circuit");
+
+        // preparing to interleave absorptions and squeezings
+        let n_iterations = max(absorptions.len(), squeezings.len());
+        absorptions.resize(n_iterations, Vec::new());
+        squeezings.resize(n_iterations, 0);
+
+        for (absorption, squeezing) in zip(absorptions, squeezings) {
+            // absorb (if any elements were provided)
+            native_sponge.update(&absorption);
+            circuit_sponge.update(&ctx.assign_witnesses(absorption));
+
+            // squeeze (if any elements were requested)
+            for _ in 0..squeezing {
+                let native_squeezed = native_sponge.squeeze();
+                let circuit_squeezed =
+                    circuit_sponge.squeeze(&mut ctx, &gate).expect("Failed to squeeze");
+
+                assert_eq!(native_squeezed, *circuit_squeezed.value());
+            }
+        }
+
+        // even if no squeezings were requested, we squeeze to verify the
+        // states are the same after all absorptions
+        let native_squeezed = native_sponge.squeeze();
+        let circuit_squeezed = circuit_sponge.squeeze(&mut ctx, &gate).expect("Failed to squeeze");
+
+        assert_eq!(native_squeezed, *circuit_squeezed.value());
+    }
+
+    fn random_nested_list_f<F: ScalarField>(len: usize, max_sub_len: usize) -> Vec<Vec<F>> {
+        let mut rng = rand::thread_rng();
+        let mut list = Vec::new();
+        for _ in 0..len {
+            let len = rng.gen_range(0..=max_sub_len);
+            let mut sublist = Vec::new();
+
+            for _ in 0..len {
+                sublist.push(F::random(&mut rng));
+            }
+            list.push(sublist);
+        }
+        list
+    }
+
+    fn random_list_usize(len: usize, max: usize) -> Vec<usize> {
+        let mut rng = rand::thread_rng();
+        let mut list = Vec::new();
+        for _ in 0..len {
+            list.push(rng.gen_range(0..=max));
+        }
+        list
+    }
+
+    #[test]
+    fn test_poseidon_compatibility_squeezing_only() {
+        let absorptions = Vec::new();
+        let squeezings = random_list_usize(10, 7);
+
+        poseidon_compatiblity_verification::<Fr, 3, 2>(absorptions, squeezings, 8, 57);
+    }
+
+    #[test]
+    fn test_poseidon_compatibility_absorbing_only() {
+        let absorptions = random_nested_list_f(8, 5);
+        let squeezings = Vec::new();
+
+        poseidon_compatiblity_verification::<Fr, 3, 2>(absorptions, squeezings, 8, 57);
+    }
+
+    #[test]
+    fn test_poseidon_compatibility_interleaved() {
+        let absorptions = random_nested_list_f(10, 5);
+        let squeezings = random_list_usize(7, 10);
+
+        poseidon_compatiblity_verification::<Fr, 3, 2>(absorptions, squeezings, 8, 57);
+    }
+
+    #[test]
+    fn test_poseidon_compatibility_other_params() {
+        let absorptions = random_nested_list_f(10, 10);
+        let squeezings = random_list_usize(10, 10);
+
+        poseidon_compatiblity_verification::<Fr, 5, 4>(absorptions, squeezings, 8, 120);
+    }
+}