chore: remove extra allocations in merge_many

0xPolygonMiden · Oct 29, 2024 · 9512741 · 9512741
1 parent 7f46e21
commit 9512741
Show file tree

Hide file tree

Showing 5 changed files with 32 additions and 19 deletions.
diff --git a/src/hash/blake/mod.rs b/src/hash/blake/mod.rs
@@ -2,7 +2,7 @@ use alloc::{string::String, vec::Vec};
 use core::{
     mem::{size_of, transmute, transmute_copy},
     ops::Deref,
-    slice::from_raw_parts,
+    slice::{self, from_raw_parts},
 };
 
 use super::{Digest, ElementHasher, Felt, FieldElement, Hasher};
@@ -33,6 +33,14 @@ const DIGEST20_BYTES: usize = 20;
 #[cfg_attr(feature = "serde", serde(into = "String", try_from = "&str"))]
 pub struct Blake3Digest<const N: usize>([u8; N]);
 
+impl<const N: usize> Blake3Digest<N> {
+    pub fn digests_as_bytes(digests: &[Blake3Digest<N>]) -> &[u8] {
+        let p = digests.as_ptr();
+        let len = digests.len() * N;
+        unsafe { slice::from_raw_parts(p as *const u8, len) }
+    }
+}
+
 impl<const N: usize> Default for Blake3Digest<N> {
     fn default() -> Self {
         Self([0; N])
@@ -115,8 +123,7 @@ impl Hasher for Blake3_256 {
     }
 
     fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        let bytes: Vec<u8> = values.iter().flat_map(|v| v.as_bytes()).collect();
-        Blake3Digest(blake3::hash(&bytes).into())
+        Blake3Digest(blake3::hash(Blake3Digest::digests_as_bytes(values)).into())
     }
 
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {

diff --git a/src/hash/rescue/rpo/digest.rs b/src/hash/rescue/rpo/digest.rs
@@ -1,5 +1,5 @@
 use alloc::string::String;
-use core::{cmp::Ordering, fmt::Display, ops::Deref};
+use core::{cmp::Ordering, fmt::Display, ops::Deref, slice};
 
 use super::{Digest, Felt, StarkField, DIGEST_BYTES, DIGEST_SIZE, ZERO};
 use crate::{
@@ -34,13 +34,19 @@ impl RpoDigest {
         <Self as Digest>::as_bytes(self)
     }
 
-    pub fn digests_as_elements<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
+    pub fn digests_as_elements_iter<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
     where
         I: Iterator<Item = &'a Self>,
     {
         digests.flat_map(|d| d.0.iter())
     }
 
+    pub fn digests_as_elements(digests: &[Self]) -> &[Felt] {
+        let p = digests.as_ptr();
+        let len = digests.len() * DIGEST_SIZE;
+        unsafe { slice::from_raw_parts(p as *const Felt, len) }
+    }
+
     /// Returns hexadecimal representation of this digest prefixed with `0x`.
     pub fn to_hex(&self) -> String {
         bytes_to_hex_string(self.as_bytes())

diff --git a/src/hash/rescue/rpo/mod.rs b/src/hash/rescue/rpo/mod.rs
@@ -1,4 +1,3 @@
-use alloc::vec::Vec;
 use core::ops::Range;
 
 use super::{
@@ -155,7 +154,7 @@ impl Hasher for Rpo256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = Self::Digest::digests_as_elements(values.iter());
+        let it = Self::Digest::digests_as_elements_iter(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }
@@ -166,9 +165,7 @@ impl Hasher for Rpo256 {
     }
 
     fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        let elements: Vec<Felt> =
-            Self::Digest::digests_as_elements(values.iter()).copied().collect();
-        Self::hash_elements(&elements)
+        Self::hash_elements(Self::Digest::digests_as_elements(values))
     }
 
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
@@ -297,7 +294,7 @@ impl Rpo256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = RpoDigest::digests_as_elements(values.iter());
+        let it = RpoDigest::digests_as_elements_iter(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }

diff --git a/src/hash/rescue/rpx/digest.rs b/src/hash/rescue/rpx/digest.rs
@@ -1,5 +1,5 @@
 use alloc::string::String;
-use core::{cmp::Ordering, fmt::Display, ops::Deref};
+use core::{cmp::Ordering, fmt::Display, ops::Deref, slice};
 
 use super::{Digest, Felt, StarkField, DIGEST_BYTES, DIGEST_SIZE, ZERO};
 use crate::{
@@ -34,13 +34,19 @@ impl RpxDigest {
         <Self as Digest>::as_bytes(self)
     }
 
-    pub fn digests_as_elements<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
+    pub fn digests_as_elements_iter<'a, I>(digests: I) -> impl Iterator<Item = &'a Felt>
     where
         I: Iterator<Item = &'a Self>,
     {
         digests.flat_map(|d| d.0.iter())
     }
 
+    pub fn digests_as_elements(digests: &[Self]) -> &[Felt] {
+        let p = digests.as_ptr();
+        let len = digests.len() * DIGEST_SIZE;
+        unsafe { slice::from_raw_parts(p as *const Felt, len) }
+    }
+
     /// Returns hexadecimal representation of this digest prefixed with `0x`.
     pub fn to_hex(&self) -> String {
         bytes_to_hex_string(self.as_bytes())

diff --git a/src/hash/rescue/rpx/mod.rs b/src/hash/rescue/rpx/mod.rs
@@ -1,4 +1,3 @@
-use alloc::vec::Vec;
 use core::ops::Range;
 
 use super::{
@@ -161,7 +160,7 @@ impl Hasher for Rpx256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = Self::Digest::digests_as_elements(values.iter());
+        let it = Self::Digest::digests_as_elements_iter(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }
@@ -172,9 +171,7 @@ impl Hasher for Rpx256 {
     }
 
     fn merge_many(values: &[Self::Digest]) -> Self::Digest {
-        let elements: Vec<Felt> =
-            Self::Digest::digests_as_elements(values.iter()).copied().collect();
-        Self::hash_elements(&elements)
+        Self::hash_elements(Self::Digest::digests_as_elements(values))
     }
 
     fn merge_with_int(seed: Self::Digest, value: u64) -> Self::Digest {
@@ -300,7 +297,7 @@ impl Rpx256 {
         // initialize the state by copying the digest elements into the rate portion of the state
         // (8 total elements), and set the capacity elements to 0.
         let mut state = [ZERO; STATE_WIDTH];
-        let it = RpxDigest::digests_as_elements(values.iter());
+        let it = RpxDigest::digests_as_elements_iter(values.iter());
         for (i, v) in it.enumerate() {
             state[RATE_RANGE.start + i] = *v;
         }