Merge pull request mit-dci#41 from Davidson-Souza/use-priority-queue

Perf: Fix a bottleneck in `calculate_hashes`
Davidson-Souza · Sep 22, 2023 · b590e73 · b590e73
2 parents 2b865d3 + 7618592
commit b590e73
Showing 1 changed file with 26 additions and 3 deletions.
diff --git a/src/accumulator/proof.rs b/src/accumulator/proof.rs
@@ -679,8 +679,10 @@ impl Proof {
         Ok(new_positions)
     }
     fn sorted_push(nodes: &mut Vec<(u64, NodeHash)>, to_add: (u64, NodeHash)) {
-        nodes.push(to_add);
-        nodes.sort();
+        let pos = nodes
+            .binary_search_by(|(pos, _)| pos.cmp(&to_add.0))
+            .unwrap_or_else(|x| x);
+        nodes.insert(pos, to_add);
     }
 }
 
@@ -1164,10 +1166,31 @@ mod tests {
 #[cfg(bench)]
 mod bench {
     extern crate test;
-    use super::Stump;
+    use crate::accumulator::stump::Stump;
     use crate::accumulator::{proof::Proof, util::hash_from_u8};
     use test::Bencher;
+    #[bench]
+    fn bench_calculate_hashes(bencher: &mut Bencher) {
+        let preimages = 0..255_u8;
+        let utxos = preimages
+            .into_iter()
+            .map(|preimage| hash_from_u8(preimage))
+            .collect::<Vec<_>>();
 
+        let (stump, modified) = Stump::new().modify(&utxos, &[], &Proof::default()).unwrap();
+        let proof = Proof::default();
+        let (proof, cached_hashes) = proof
+            .update(
+                vec![],
+                utxos.clone(),
+                vec![],
+                (0..128).into_iter().collect(),
+                modified,
+            )
+            .unwrap();
+
+        bencher.iter(|| proof.calculate_hashes(&cached_hashes, stump.leaves))
+    }
     #[bench]
     fn bench_proof_update(bencher: &mut Bencher) {
         let preimages = [0_u8, 1, 2, 3, 4, 5];