From 27a42cddca35c677594d7cb2d295917ade3a963f Mon Sep 17 00:00:00 2001
From: Remco Bloemen <remco@wicked.ventures>
Date: Wed, 7 May 2025 18:15:27 +0200
Subject: [PATCH] Fix merkle hash bench

---
 merkle-hash-bench/src/hashes/plonky3.rs            | 14 +++++++++-----
 merkle-hash-bench/src/hashes/poseidon2_t2_ruint.rs |  6 +++---
 merkle-hash-bench/src/hashes/poseidon2_t3_ruint.rs |  6 +++---
 merkle-hash-bench/src/main.rs                      |  2 +-
 merkle-hash-bench/src/mod_ring/ruint.rs            |  6 +++---
 5 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/merkle-hash-bench/src/hashes/plonky3.rs b/merkle-hash-bench/src/hashes/plonky3.rs
index f60383d64..7c59b5d8a 100644
--- a/merkle-hash-bench/src/hashes/plonky3.rs
+++ b/merkle-hash-bench/src/hashes/plonky3.rs
@@ -1,6 +1,6 @@
 #![cfg(feature = "plonky3")]
 use {
-    crate::{register_hash, Field, HashFn, SmolHasher},
+    crate::{mod_ring::fields::Bn254Field, register_hash, Field, HashFn, SmolHasher},
     bytemuck::cast_slice_mut,
     p3_bn254_fr::{Bn254Fr, FFBn254Fr, Poseidon2Bn254},
     p3_field::{
@@ -12,6 +12,7 @@ use {
     p3_rescue::Rescue,
     p3_symmetric::Permutation,
     rand::rng,
+    std::mem::transmute,
 };
 
 type RescueGoldilocks = Rescue<Goldilocks, MdsMatrixGoldilocks, 8, 7>;
@@ -105,7 +106,7 @@ impl SmolHasher for Poseidon2Bn254<3> {
                 Bn254Fr::ZERO,
             ];
             let state = self.permute(state);
-            hash.copy_from_slice(state[0].value.to_bytes().as_slice());
+            hash.copy_from_slice(bytes_from_fr(state[0]).as_slice());
         }
     }
 }
@@ -113,9 +114,12 @@ impl SmolHasher for Poseidon2Bn254<3> {
 fn fr_from_bytes(bytes: &[u8]) -> Bn254Fr {
     let mut bytes: [u8; 32] = bytes.try_into().unwrap();
     bytes[31] = 0; // Force smaller than modulus.
-    Bn254Fr {
-        value: FFBn254Fr::from_bytes(&bytes).unwrap(),
-    }
+    let element = FFBn254Fr::from_bytes(&bytes).unwrap();
+    unsafe { transmute(element) }
+}
+
+fn bytes_from_fr(element: Bn254Fr) -> [u8; 32] {
+    unsafe { transmute(element) }
 }
 
 impl SmolHasher for MonolithMersenne31<MdsMatrixMersenne31, 16, 5> {
diff --git a/merkle-hash-bench/src/hashes/poseidon2_t2_ruint.rs b/merkle-hash-bench/src/hashes/poseidon2_t2_ruint.rs
index 4b29c840a..c86ea765d 100644
--- a/merkle-hash-bench/src/hashes/poseidon2_t2_ruint.rs
+++ b/merkle-hash-bench/src/hashes/poseidon2_t2_ruint.rs
@@ -54,9 +54,9 @@ impl Poseidon2T2Ruint {
     pub fn new() -> Self {
         let mut rng = rand::rng();
         Self {
-            first:  rng.gen(),
-            middle: array::from_fn(|_| rng.gen()),
-            last:   rng.gen(),
+            first:  rng.random(),
+            middle: array::from_fn(|_| rng.random()),
+            last:   rng.random(),
         }
     }
 
diff --git a/merkle-hash-bench/src/hashes/poseidon2_t3_ruint.rs b/merkle-hash-bench/src/hashes/poseidon2_t3_ruint.rs
index caa8f1c09..f940076ad 100644
--- a/merkle-hash-bench/src/hashes/poseidon2_t3_ruint.rs
+++ b/merkle-hash-bench/src/hashes/poseidon2_t3_ruint.rs
@@ -57,9 +57,9 @@ impl Poseidon2T3Ruint {
     pub fn new() -> Self {
         let mut rng = rand::rng();
         Self {
-            first:  rng.gen(),
-            middle: array::from_fn(|_| rng.gen()),
-            last:   rng.gen(),
+            first:  rng.random(),
+            middle: array::from_fn(|_| rng.random()),
+            last:   rng.random(),
         }
     }
 
diff --git a/merkle-hash-bench/src/main.rs b/merkle-hash-bench/src/main.rs
index 3716d06e6..53f9dd4fb 100644
--- a/merkle-hash-bench/src/main.rs
+++ b/merkle-hash-bench/src/main.rs
@@ -131,7 +131,7 @@ struct Args {
 
 fn main() -> Result<()> {
     let args: Args = argh::from_env();
-    let duration = Duration::from_secs_f64(args.duration.unwrap_or(0.01));
+    let duration = Duration::from_secs_f64(args.duration.unwrap_or(0.1));
 
     // Consrtuct all hashers.
     let hashes = HASHES.iter().map(|ctor| ctor()).collect::<Vec<_>>();
diff --git a/merkle-hash-bench/src/mod_ring/ruint.rs b/merkle-hash-bench/src/mod_ring/ruint.rs
index 0751b6946..82838f79a 100644
--- a/merkle-hash-bench/src/mod_ring/ruint.rs
+++ b/merkle-hash-bench/src/mod_ring/ruint.rs
@@ -28,9 +28,9 @@ impl<const BITS: usize, const LIMBS: usize> UintMont for Uint<BITS, LIMBS> {
 
     fn random<R: rand::Rng + ?Sized>(rng: &mut R, max: Self) -> Self {
         let mut result = Self::ZERO;
-        unsafe { // TODO: Update rand crate,
-            for limb in  result.as_limbs_mut() {
-                *limb = rng.gen();
+        unsafe {
+            for limb in result.as_limbs_mut() {
+                *limb = rng.random();
             }
         }
         result %= max;