feat: optimize variable length hashes (#32)

Author: jialinli98

src/poseidon2.nr

@@ -69,23 +69,40 @@ impl Poseidon2 {
                 state[2] += input[i * RATE + 2];
                 state = crate::poseidon2_permutation(state, 4);
             }
-            // handle remaining (<3) elements
+
+            // handle remaining elements after last full RATE-sized chunk
             let remainder_start = (in_len / RATE) * RATE;
             for j in remainder_start..in_len {
                 state[j - remainder_start] += input[j];
             }
         } else {
-            for i in 0..input.len() {
-                if i < in_len {
-                    state[i % RATE] += input[i];
-                    if (i + 1) % RATE == 0 {
-                        state = crate::poseidon2_permutation(state, 4);
-                    }
+            let mut states: [[Field; 4]; N / RATE + 1] = [[0; 4]; N / RATE + 1];
+            states[0] = state;
+
+            // process all full RATE-sized chunks, storing state after each permutation
+            for chunk_idx in 0..(N / RATE) {
+                for i in 0..RATE {
+                    state[i] += input[chunk_idx * RATE + i];
+                }
+                state = crate::poseidon2_permutation(state, 4);
+                states[chunk_idx + 1] = state;
+            }
+
+            // get state at the last full block before in_len
+            let first_partially_filled_chunk = in_len / RATE;
+            state = states[first_partially_filled_chunk];
+
+            // handle remaining elements after last full RATE-sized chunk
+            let remainder_start = (in_len / RATE) * RATE;
+            for j in 0..RATE {
+                let idx = remainder_start + j;
+                if idx < in_len {
+                    state[j] += input[idx];
                 }
             }
         }
 
-        // Always run final permutation unless we just completed a full chunk
+        // always run final permutation unless we just completed a full chunk
         // still need to permute once if in_len is 0
         if (in_len == 0) | (in_len % RATE != 0) {
             state = crate::poseidon2_permutation(state, 4)

src/tests.nr

@@ -1,4 +1,4 @@
-use super::poseidon;
+use super::{poseidon, poseidon2};
 
 #[test]
 fn reference_impl_test_vectors() {
@@ -65,3 +65,27 @@ fn test_poseidon_hash_8() {
     let expected = 0x2921ab9bd0140cbc98e40395c0fefb40337a4d54fbbecd9a4d43b3d8d0c4d8d1;
     assert_eq(hash, expected);
 }
+
+#[test]
+fn test_poseidon2_1() {
+    let input = [1000];
+    let hash = poseidon2::Poseidon2::hash(input, input.len());
+    let expected = 0x16433a80e26a23547e25d61dd95fd5793d1ca2dcd78ae64cd146d3b99a35fa7c;
+    assert_eq(hash, expected);
+}
+
+#[test]
+fn test_poseidon2_2() {
+    let input = [1000, 2000];
+    let hash = poseidon2::Poseidon2::hash(input, input.len());
+    let expected = 0x118d5a5ecb25dafe99eb45cb196604a23d0b7c0cbd0c2be29e0787e59b7a1d8a;
+    assert_eq(hash, expected);
+}
+
+#[test]
+fn test_poseidon2_3() {
+    let input = [1000, 2000, 3000];
+    let hash = poseidon2::Poseidon2::hash(input, input.len());
+    let expected = 0x0f1badcd0d52ced816fb6e6826fdf66ada038135d53cbb993f320ca6529223cd;
+    assert_eq(hash, expected);
+}