perf: refactor merkle update

broxus · Nov 29, 2024 · d13f4fe · d13f4fe
1 parent bea5ae3
commit d13f4fe
Show file tree

Hide file tree

Showing 4 changed files with 126 additions and 74 deletions.
diff --git a/src/cell/usage_tree.rs b/src/cell/usage_tree.rs
@@ -102,6 +102,10 @@ impl UsageTreeWithSubtrees {
     pub fn add_subtree(&mut self, root: &DynCell) -> bool {
         self.subtrees.insert(*root.repr_hash())
     }
+
+    pub(crate) fn len(&self) -> usize {
+        self.state.len()
+    }
 }
 
 #[cfg(not(feature = "sync"))]

diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs
@@ -30,6 +30,9 @@ mod __checks {
 pub trait MerkleFilter {
     /// Returns how the cell should be included in the Merkle proof or update.
     fn check(&self, cell: &HashBytes) -> FilterAction;
+
+    /// Returns the number of elements in the filter, if known.
+    fn size_hint(&self) -> Option<usize>;
 }
 
 /// Merkle filter action.
@@ -48,6 +51,11 @@ impl<T: MerkleFilter + ?Sized> MerkleFilter for &T {
     fn check(&self, cell: &HashBytes) -> FilterAction {
         <T as MerkleFilter>::check(self, cell)
     }
+
+    #[inline]
+    fn size_hint(&self) -> Option<usize> {
+        <T as MerkleFilter>::size_hint(self)
+    }
 }
 
 impl MerkleFilter for UsageTree {
@@ -58,6 +66,10 @@ impl MerkleFilter for UsageTree {
             FilterAction::Skip
         }
     }
+
+    fn size_hint(&self) -> Option<usize> {
+        Some(UsageTree::len(self))
+    }
 }
 
 impl MerkleFilter for UsageTreeWithSubtrees {
@@ -70,6 +82,10 @@ impl MerkleFilter for UsageTreeWithSubtrees {
             FilterAction::Skip
         }
     }
+
+    fn size_hint(&self) -> Option<usize> {
+        Some(self.len())
+    }
 }
 
 impl<S: BuildHasher> MerkleFilter for HashSet<HashBytes, S> {
@@ -80,6 +96,10 @@ impl<S: BuildHasher> MerkleFilter for HashSet<HashBytes, S> {
             FilterAction::Skip
         }
     }
+
+    fn size_hint(&self) -> Option<usize> {
+        Some(self.len())
+    }
 }
 
 impl<S: BuildHasher> MerkleFilter for HashSet<&HashBytes, S> {
@@ -90,4 +110,8 @@ impl<S: BuildHasher> MerkleFilter for HashSet<&HashBytes, S> {
             FilterAction::Skip
         }
     }
+
+    fn size_hint(&self) -> Option<usize> {
+        Some(self.len())
+    }
 }
diff --git a/src/merkle/proof.rs b/src/merkle/proof.rs
@@ -190,6 +190,10 @@ impl MerkleProof {
                     FilterAction::Skip
                 }
             }
+
+            fn size_hint(&self) -> Option<usize> {
+                Some(self.cells.len())
+            }
         }
 
         let mut stack = vec![root.references()];

diff --git a/src/merkle/update.rs b/src/merkle/update.rs
@@ -486,78 +486,6 @@ struct BuilderImpl<'a, 'b> {
 
 impl<'a: 'b, 'b> BuilderImpl<'a, 'b> {
     fn build(self) -> Result<MerkleUpdate, Error> {
-        struct Resolver<'a, S> {
-            pruned_branches: HashMap<&'a HashBytes, bool, S>,
-            visited: HashSet<&'a HashBytes, S>,
-            filter: &'a dyn MerkleFilter,
-            changed_cells: HashSet<&'a HashBytes, S>,
-        }
-
-        impl<'a, S> Resolver<'a, S>
-        where
-            S: BuildHasher,
-        {
-            fn fill(&mut self, cell: &'a DynCell, mut skip_filter: bool) -> bool {
-                let repr_hash = cell.repr_hash();
-
-                // Skip visited cells
-                if self.visited.contains(repr_hash) {
-                    return false;
-                }
-                self.visited.insert(repr_hash);
-
-                let is_pruned = match self.pruned_branches.get_mut(repr_hash) {
-                    Some(true) => return false,
-                    Some(visited) => {
-                        *visited = true;
-                        true
-                    }
-                    None => false,
-                };
-
-                let process_children = if skip_filter {
-                    true
-                } else {
-                    match self.filter.check(repr_hash) {
-                        FilterAction::Skip => false,
-                        FilterAction::Include => true,
-                        FilterAction::IncludeSubtree => {
-                            skip_filter = true;
-                            true
-                        }
-                    }
-                };
-
-                let mut result = false;
-                if process_children {
-                    for child in cell.references() {
-                        result |= self.fill(child, skip_filter);
-                    }
-
-                    if result {
-                        self.changed_cells.insert(repr_hash);
-                    }
-                }
-
-                result | is_pruned
-            }
-        }
-
-        struct InvertedFilter<F>(F);
-
-        impl<F: MerkleFilter> MerkleFilter for InvertedFilter<F> {
-            #[inline]
-            fn check(&self, cell: &HashBytes) -> FilterAction {
-                if self.0.check(cell) == FilterAction::Skip {
-                    // TODO: check if FilterAction::IncludeSubtree is correct,
-                    // because it is more optimal to just include the new subtree
-                    FilterAction::Include
-                } else {
-                    FilterAction::Skip
-                }
-            }
-        }
-
         let old_hash = self.old.repr_hash();
         let old_depth = self.old.repr_depth();
         let new_hash = self.new.repr_hash();
@@ -590,9 +518,15 @@ impl<'a: 'b, 'b> BuilderImpl<'a, 'b> {
         // Prepare cell diff resolver
         let mut resolver = Resolver {
             pruned_branches,
-            visited: Default::default(),
+            visited: HashSet::with_capacity_and_hasher(
+                self.filter.size_hint().unwrap_or(512),
+                Default::default(),
+            ),
             filter: self.filter,
-            changed_cells: Default::default(),
+            changed_cells: HashSet::with_capacity_and_hasher(
+                self.filter.size_hint().unwrap_or(512),
+                Default::default(),
+            ),
         };
 
         // Find all changed cells in the old cell tree
@@ -619,6 +553,92 @@ impl<'a: 'b, 'b> BuilderImpl<'a, 'b> {
     }
 }
 
+struct Resolver<'a, S> {
+    pruned_branches: HashMap<&'a HashBytes, bool, S>,
+    visited: HashSet<&'a HashBytes, S>,
+    filter: &'a dyn MerkleFilter,
+    changed_cells: HashSet<&'a HashBytes, S>,
+}
+
+impl<'a, S> Resolver<'a, S>
+where
+    S: BuildHasher,
+{
+    fn fill(&mut self, cell: &'a DynCell, mut skip_filter: bool) -> bool {
+        let repr_hash = cell.repr_hash();
+
+        if !self.visited.insert(repr_hash) {
+            return false;
+        }
+
+        let is_pruned = match self.pruned_branches.get_mut(repr_hash) {
+            Some(visited) => {
+                if *visited {
+                    return false;
+                }
+                *visited = true;
+                true
+            }
+            None => false,
+        };
+
+        let (process_children, new_skip_filter) = if skip_filter {
+            (true, true)
+        } else {
+            match self.filter.check(repr_hash) {
+                FilterAction::Skip => (false, false),
+                FilterAction::Include => (true, false),
+                FilterAction::IncludeSubtree => (true, true),
+            }
+        };
+        skip_filter = new_skip_filter;
+
+        // Process children only if needed
+        if !process_children {
+            return is_pruned;
+        }
+
+        let mut result = false;
+        let refs = cell.references();
+
+        // Process all references
+        for child in refs {
+            result |= self.fill(child, skip_filter);
+
+            // early exit if we found changes and don't need to process all children
+            if result && !skip_filter {
+                break;
+            }
+        }
+
+        // update changed cells if needed
+        if result {
+            self.changed_cells.insert(repr_hash);
+        }
+
+        result | is_pruned
+    }
+}
+
+struct InvertedFilter<F>(F);
+
+impl<F: MerkleFilter> MerkleFilter for InvertedFilter<F> {
+    #[inline]
+    fn check(&self, cell: &HashBytes) -> FilterAction {
+        if self.0.check(cell) == FilterAction::Skip {
+            // TODO: check if FilterAction::IncludeSubtree is correct,
+            // because it is more optimal to just include the new subtree
+            FilterAction::Include
+        } else {
+            FilterAction::Skip
+        }
+    }
+
+    fn size_hint(&self) -> Option<usize> {
+        self.0.size_hint()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;