rhp: use hash_many for sector_root

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()

src/rhp.rs

@@ -1,5 +1,6 @@
-use crate::merkle::{sum_leaf, sum_node};
+use crate::merkle::{sum_node, LEAF_HASH_PREFIX, NODE_HASH_PREFIX};
 use crate::Hash256;
+use blake2b_simd::many::{hash_many, HashManyJob};
 use blake2b_simd::Params;
 use rayon::prelude::*;
 
@@ -12,21 +13,75 @@ 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] = LEAF_HASH_PREFIX[0];
+                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 = 8;
     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]);
-            });
+        tree_hashes.par_chunks_mut(chunk_size).for_each(|nodes| {
+            // prepare inputs
+            let mut inputs = [[0u8; 65]; 4];
+            for (j, input) in inputs.iter_mut().enumerate() {
+                input[0] = NODE_HASH_PREFIX[0];
+                let step = j * chunk_size / 4;
+                input[1..33].copy_from_slice(&nodes[step]);
+                input[33..65].copy_from_slice(&nodes[step + chunk_size / 8]);
+            }
+
+            // 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
+            nodes[0] = jobs[0].to_hash().as_bytes().try_into().unwrap();
+            nodes[nodes.len() / 4] = jobs[1].to_hash().as_bytes().try_into().unwrap();
+            nodes[nodes.len() / 4 * 2] = jobs[2].to_hash().as_bytes().try_into().unwrap();
+            nodes[nodes.len() / 4 * 3] = jobs[3].to_hash().as_bytes().try_into().unwrap();
+        });
         chunk_size *= 2;
     }
-    Hash256::from(tree_hashes[0])
+    // hash last 4 nodes manually
+    let left = sum_node(
+        &params,
+        &tree_hashes[0],
+        &tree_hashes[tree_hashes.len() / 4],
+    );
+    let right = sum_node(
+        &params,
+        &tree_hashes[tree_hashes.len() / 4 * 2],
+        &tree_hashes[tree_hashes.len() / 4 * 3],
+    );
+    Hash256::from(sum_node(&params, &left, &right))
 }
 
 #[cfg(test)]