optimize Poseidon for larger number of rounds

Author: shamatar

Cargo.toml

@@ -25,8 +25,8 @@ serde_derive = "1.0.80"
 tiny-keccak = "1.4.2"
 rust-crypto = "0.2"
 
-#bellman_ce = { path = "../bellman"}
-bellman_ce = { version = "0.3.0", default-features = false}
+bellman_ce = { path = "../bellman"}
+#bellman_ce = { version = "0.3.0", default-features = false}
 blake2-rfc_bellman_edition = "0.0.1"
 
 [dev-dependencies]

src/circuit/num.rs

@@ -778,9 +778,30 @@ impl<E: Engine> Num<E> {
 
         self.value = newval;
         let mut lc = LinearCombination::zero();
-        std::mem::swap(&mut self.lc, &mut lc);
+        // std::mem::swap(&mut self.lc, &mut lc);
+        use std::collections::HashMap;
+        let mut final_coeffs: HashMap<bellman::Variable, E::Fr> = HashMap::new();
+        for (var, coeff) in self.lc.as_ref() {
+            if final_coeffs.get(var).is_some() {
+                if let Some(existing_coeff) = final_coeffs.get_mut(var) {
+                    existing_coeff.add_assign(&coeff);
+                } 
+            } else {
+                final_coeffs.insert(*var, *coeff);
+            }
+        }
+
         for (var, coeff) in other.lc.as_ref() {
-            lc = lc + (*coeff, var.clone());
+            if final_coeffs.get(var).is_some() {
+                if let Some(existing_coeff) = final_coeffs.get_mut(var) {
+                    existing_coeff.add_assign(&coeff);
+                }
+            } else {
+                final_coeffs.insert(*var, *coeff);
+            }
+        }
+        for (var, coeff) in final_coeffs.into_iter() {
+            lc = lc + (coeff, var);
         }
         self.lc = lc;
     }

src/circuit/poseidon_hash.rs

@@ -100,11 +100,6 @@ impl<E: PoseidonEngine> QuinticSBox<E> {
     }
 }
 
-// pub fn poseidon_tree_hash<E: PoseidonEngine<SBox = QuinticSBox<E> >, CS>(
-//     mut cs: CS,
-//     elements: &[AllocatedNum<E>],
-//     params: &E::Params
-
 pub fn poseidon_hash<E: PoseidonEngine<SBox = QuinticSBox<E> >, CS>(
     mut cs: CS,
     input: &[AllocatedNum<E>],
@@ -311,6 +306,12 @@ fn poseidon_mimc_round<E: PoseidonEngine<SBox = QuinticSBox<E> >, CS>(
         round += 1;
     }
 
+    // up to this point we always made a well-formed LC that later was collapsed into
+    // a signel variable after non-linearity application
+    // now we need to make linear combinations of linear combinations, so basically make
+    // filtering and joining. It's actually possible to just separate MSD matrix into
+    // three in later optimizations
+
     // now we need to apply full SBox of the last full round, then do linear
     // transformation and add first round constants before going through partial rounds
     {
@@ -329,6 +330,9 @@ fn poseidon_mimc_round<E: PoseidonEngine<SBox = QuinticSBox<E> >, CS>(
         add_round_constants::<E, CS>(params, &mut linear_transformation_results[..], 0, false);
         state = linear_transformation_results;
 
+        // up to this point linear combinations are well-formed and have number
+        // of terms equal to the number of variables in the state
+
         round += 1;
     }
 
@@ -462,7 +466,7 @@ mod test {
         let mut rng = XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);
         let params = Bn256PoseidonParams::new::<BlakeHasher>();
         let input: Vec<Fr> = (0..params.t()).map(|_| rng.gen()).collect();
-        let expected = poseidon::poseidon_hash::<Bn256>(&params, &input[..]);
+        let expected = poseidon::poseidon_mimc::<Bn256>(&params, &input[..]);
 
         {
             let mut cs = TestConstraintSystem::<Bn256>::new();
@@ -482,7 +486,7 @@ mod test {
             ).unwrap();
 
             assert!(cs.is_satisfied());
-            assert!(res.len() == 1);
+            assert!(res.len() == (params.t() as usize));
 
             assert_eq!(res[0].get_value().unwrap(), expected[0]);
         }