From 725d6397b8572d3fc7b96bdcb674475550ecaa96 Mon Sep 17 00:00:00 2001
From: arnaucube <root@arnaucube.com>
Date: Wed, 22 Jul 2020 19:53:00 +0200
Subject: [PATCH] Update methods to more efficient usage of ff

On a Intel(R) Core(TM) i5-6300U CPU @ 2.40GHz, with 16GB of RAM:
- Old (using num-bigint):
hash  time:   [5.9258 ms 5.9407 ms 5.9587 ms]
- New (using ff):
hash  time:   [120.12 us 121.08 us 122.30 us]

On a Intel(R) Core(TM) i7-8705G CPU @ 3.10GHz, with 32 GB of RAM:
- Old (using num-bigint):
hash  time:   [4.1192 ms 4.1318 ms 4.1461 ms]
- In the previous commit (using ff):
hash  time:   [91.394 us 91.430 us 91.476 us]
- In this commit (using ff):
hash  time:   [85.517 us 85.545 us 85.574 us]
---
 README.md                      |  4 ----
 benches/bench_poseidon_hash.rs |  2 +-
 src/lib.rs                     | 44 ++++++++++++++++++----------------
 3 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/README.md b/README.md
index 651d35c..3b0ba06 100644
--- a/README.md
+++ b/README.md
@@ -8,7 +8,3 @@ Compatible with the Poseidon Go implementation done in https://github.com/iden3/
 ## Warning
 Do not use in production
 
-## Benchmarks
-```
-hash           time:   [120.12 us 121.08 us 122.30 us]
-```
diff --git a/benches/bench_poseidon_hash.rs b/benches/bench_poseidon_hash.rs
index 6420ce2..958acaf 100644
--- a/benches/bench_poseidon_hash.rs
+++ b/benches/bench_poseidon_hash.rs
@@ -1,4 +1,4 @@
-use criterion::{black_box, criterion_group, criterion_main, Criterion};
+use criterion::{criterion_group, criterion_main, Criterion};
 
 extern crate rand;
 #[macro_use]
diff --git a/src/lib.rs b/src/lib.rs
index 2953210..a79480c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -164,32 +164,26 @@ impl Poseidon {
             constants: load_constants(),
         }
     }
-    pub fn ark(&self, state: &Vec<Fr>, c: &Fr) -> Vec<Fr> {
-        let mut new_state: Vec<Fr> = state.clone();
+    pub fn ark(&self, state: &mut Vec<Fr>, c: &Fr) {
         for i in 0..state.len() {
-            new_state[i] = state[i];
-            new_state[i].add_assign(c);
+            state[i].add_assign(c);
         }
-
-        new_state
     }
 
-    pub fn sbox(&self, state: &Vec<Fr>, i: usize) -> Vec<Fr> {
-        let mut new_state: Vec<Fr> = state.clone();
+    pub fn sbox(&self, state: &mut Vec<Fr>, i: usize) {
         if i < NROUNDSF / 2 || i >= NROUNDSF / 2 + NROUNDSP {
             for j in 0..T {
-                new_state[j] = state[j];
-                new_state[j].square();
-                new_state[j].square();
-                new_state[j].mul_assign(&state[j]);
+                let aux = state[j];
+                state[j].square();
+                state[j].square();
+                state[j].mul_assign(&aux);
             }
         } else {
-            new_state[0] = state[0];
-            new_state[0].square();
-            new_state[0].square();
-            new_state[0].mul_assign(&state[0]);
+            let aux = state[0];
+            state[0].square();
+            state[0].square();
+            state[0].mul_assign(&aux);
         }
-        new_state
     }
 
     pub fn mix(&self, state: &Vec<Fr>, m: &Vec<Vec<Fr>>) -> Vec<Fr> {
@@ -216,8 +210,8 @@ impl Poseidon {
         }
 
         for i in 0..(NROUNDSF + NROUNDSP) {
-            state = self.ark(&state, &self.constants.c[i]);
-            state = self.sbox(&state, i);
+            self.ark(&mut state, &self.constants.c[i]);
+            self.sbox(&mut state, i);
             state = self.mix(&state, &self.constants.m);
         }
 
@@ -236,13 +230,21 @@ mod tests {
             "0000000000000000000000000000000000000000000000000000000000000002",
             to_hex(&a)
         );
-        println!("`2` into hex = {}", to_hex(&a));
+
+        let b: Fr = Fr::from_str(
+            "21888242871839275222246405745257275088548364400416034343698204186575808495619",
+        )
+        .unwrap();
+        assert_eq!(
+            "0000000000000000000000000000000000000000000000000000000000000002",
+            to_hex(&b)
+        );
+        assert_eq!(&a, &b);
     }
 
     #[test]
     fn test_load_constants() {
         let constants = load_constants();
-        println!("{:?}", constants.c[0].to_string());
         assert_eq!(
             constants.c[0].to_string(),
             "Fr(0x1fd4a35e68f0946f8f5dfd2ac9d7882ce2466ec1c9766f69b5a14c3f84a17be2)"