rhp: use hash_many for sector_root

SiaFoundation · Sep 28, 2024 · bbd64c3 · bbd64c3
1 parent f659654
commit bbd64c3
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 13 deletions.
diff --git a/src/merkle.rs b/src/merkle.rs
@@ -53,6 +53,7 @@ impl Accumulator {
     }
 }
 
+#[allow(dead_code)]
 pub fn sum_leaf(params: &Params, leaf: &[u8]) -> [u8; 32] {
     let h = params
         .to_state()

diff --git a/src/rhp.rs b/src/rhp.rs
@@ -1,5 +1,6 @@
-use crate::merkle::{sum_leaf, sum_node};
+use crate::merkle::sum_node;
 use crate::Hash256;
+use blake2b_simd::many::{hash_many, HashManyJob};
 use blake2b_simd::Params;
 use rayon::prelude::*;
 
@@ -12,21 +13,65 @@ pub fn sector_root(sector: &[u8]) -> Hash256 {
     let mut params = Params::new();
     params.hash_length(32);
 
-    let mut tree_hashes = sector
-        .par_chunks_exact(SEGMENT_SIZE)
-        .map(|chunk| sum_leaf(&params, chunk))
-        .collect::<Vec<_>>();
+    let mut tree_hashes = vec![[0; 32]; SECTOR_SIZE / SEGMENT_SIZE];
+    tree_hashes
+        .par_chunks_exact_mut(4)
+        .enumerate()
+        .for_each(|(i, chunk)| {
+            // prepare inputs
+            let mut inputs = [[0u8; SEGMENT_SIZE + 1]; 4];
+            for (j, input) in inputs.iter_mut().enumerate() {
+                input[0] = 0; // prefix
+                input[1..]
+                    .copy_from_slice(&sector[SEGMENT_SIZE * (i + j)..SEGMENT_SIZE * (i + j + 1)]);
+            }
 
-    let mut chunk_size = 2;
+            // hash them
+            let mut jobs = [
+                HashManyJob::new(&params, &inputs[0][..]),
+                HashManyJob::new(&params, &inputs[1][..]),
+                HashManyJob::new(&params, &inputs[2][..]),
+                HashManyJob::new(&params, &inputs[3][..]),
+            ];
+            hash_many(&mut jobs);
+
+            // collect results
+            for j in 0..4 {
+                chunk[j] = jobs[j].to_hash().as_bytes().try_into().unwrap();
+            }
+        });
+
+    let mut chunk_size = 4;
     while chunk_size <= tree_hashes.len() {
-        tree_hashes
-            .par_chunks_exact_mut(chunk_size)
-            .for_each(|nodes| {
-                nodes[0] = sum_node(&params, &nodes[0], &nodes[nodes.len() / 2]);
-            });
-        chunk_size *= 2;
+        tree_hashes.par_chunks_mut(chunk_size).for_each(|nodes| {
+            // prepare inputs
+            let mut inputs = [[0u8; 65]; 2];
+            for (j, input) in inputs.iter_mut().enumerate() {
+                input[0] = 1; // prefix
+                let step = j * chunk_size / 2;
+                input[1..33].copy_from_slice(&nodes[step]);
+                input[33..65].copy_from_slice(&nodes[step + chunk_size / 4]);
+            }
+
+            // hash them
+            let mut jobs = [
+                HashManyJob::new(&params, &inputs[0][..]),
+                HashManyJob::new(&params, &inputs[1][..]),
+            ];
+            hash_many(&mut jobs);
+
+            // collect results
+            nodes[0] = jobs[0].to_hash().as_bytes().try_into().unwrap();
+            nodes[chunk_size / 2] = jobs[1].to_hash().as_bytes().try_into().unwrap();
+        });
+        chunk_size *= 4;
     }
-    Hash256::from(tree_hashes[0])
+    // hash last two nodes into roots
+    Hash256::from(sum_node(
+        &params,
+        &tree_hashes[0],
+        &tree_hashes[tree_hashes.len() / 2],
+    ))
 }
 
 #[cfg(test)]