From 1cb7118cf7dacb907a441e73eff69f5340edfe8d Mon Sep 17 00:00:00 2001
From: Akosh Farkash <aakoshh@gmail.com>
Date: Thu, 10 Oct 2024 13:35:04 +0100
Subject: [PATCH] Add is_variable_length to the Rust poseidon_hash function

---
 acvm-repo/bn254_blackbox_solver/src/poseidon2.rs | 14 ++++++++++++--
 tooling/nargo_cli/tests/stdlib-props.rs          |  5 ++---
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/acvm-repo/bn254_blackbox_solver/src/poseidon2.rs b/acvm-repo/bn254_blackbox_solver/src/poseidon2.rs
index 64823e37029..3aa735388ca 100644
--- a/acvm-repo/bn254_blackbox_solver/src/poseidon2.rs
+++ b/acvm-repo/bn254_blackbox_solver/src/poseidon2.rs
@@ -544,13 +544,23 @@ impl<'a> Poseidon2<'a> {
 }
 
 /// Performs a poseidon hash with a sponge construction equivalent to the one in poseidon2.nr
-pub fn poseidon_hash(inputs: &[FieldElement]) -> Result<FieldElement, BlackBoxResolutionError> {
+///
+/// The `is_variable_length` parameter is there to so we can produce an equivalent hash with
+/// the Barretenberg implementation which distinguishes between variable and fixed length inputs.
+/// Set it to true if the input length matches the static size expected by the Noir function.
+pub fn poseidon_hash(
+    inputs: &[FieldElement],
+    is_variable_length: bool,
+) -> Result<FieldElement, BlackBoxResolutionError> {
     let two_pow_64 = 18446744073709551616_u128.into();
     let iv = FieldElement::from(inputs.len()) * two_pow_64;
     let mut sponge = Poseidon2Sponge::new(iv, 3);
     for input in inputs.iter() {
         sponge.absorb(*input)?;
     }
+    if is_variable_length {
+        sponge.absorb(FieldElement::from(1u32))?;
+    }
     sponge.squeeze()
 }
 
@@ -640,7 +650,7 @@ mod test {
             FieldElement::from(3u128),
             FieldElement::from(4u128),
         ];
-        let result = super::poseidon_hash(&fields).expect("should hash successfully");
+        let result = super::poseidon_hash(&fields, false).expect("should hash successfully");
         assert_eq!(
             result,
             field_from_hex("130bf204a32cac1f0ace56c78b731aa3809f06df2731ebcf6b3464a15788b1b9"),
diff --git a/tooling/nargo_cli/tests/stdlib-props.rs b/tooling/nargo_cli/tests/stdlib-props.rs
index d6a8fdfc920..9347f596257 100644
--- a/tooling/nargo_cli/tests/stdlib-props.rs
+++ b/tooling/nargo_cli/tests/stdlib-props.rs
@@ -259,14 +259,13 @@ fn fuzz_sha512_equivalence() {
 fn fuzz_poseidon2_equivalence() {
     use bn254_blackbox_solver::poseidon_hash;
 
-    for max_len in [0, 1, 4, 511, 512] {
+    for max_len in [0, 1, 3, 4, 511, 512] {
         let source = format!(
             "fn main(input: [Field; {max_len}], message_size: u32) -> pub Field {{
                 std::hash::poseidon2::Poseidon2::hash(input, message_size)
             }}"
         );
 
-        // The actual input length can be up to the maximum.
         let strategy = (0..=max_len)
             .prop_flat_map(|len: usize| {
                 // Generate Field elements from random 32 byte vectors.
@@ -277,7 +276,7 @@ fn fuzz_poseidon2_equivalence() {
             })
             .prop_map(move |mut msg| {
                 // The output hash is a single field element.
-                let output = poseidon_hash(&msg).expect("failed to hash");
+                let output = poseidon_hash(&msg, msg.len() < max_len).expect("failed to hash");
 
                 // The input has to be padded to the maximum length.
                 let msg_size = msg.len();