From 1b0f625f1d315473b86fc38153a7c5ed27d27ffa Mon Sep 17 00:00:00 2001
From: Alexey Shekhirin <a.shekhirin@gmail.com>
Date: Mon, 28 Oct 2024 11:06:26 +0000
Subject: [PATCH] perf(trie): collect only changed sparse nodes at a depth
 (#12093)

---
 crates/trie/sparse/src/trie.rs | 47 ++++++++++++++++++++++++----------
 1 file changed, 33 insertions(+), 14 deletions(-)

diff --git a/crates/trie/sparse/src/trie.rs b/crates/trie/sparse/src/trie.rs
index e7ee66c540..91362eed52 100644
--- a/crates/trie/sparse/src/trie.rs
+++ b/crates/trie/sparse/src/trie.rs
@@ -569,27 +569,36 @@ impl RevealedSparseTrie {
     /// Update hashes of the nodes that are located at a level deeper than or equal to the provided
     /// depth. Root node has a level of 0.
     pub fn update_rlp_node_level(&mut self, depth: usize) {
-        let targets = self.get_nodes_at_depth(depth);
         let mut prefix_set = self.prefix_set.clone().freeze();
+
+        let targets = self.get_changed_nodes_at_depth(&mut prefix_set, depth);
         for target in targets {
             self.rlp_node(target, &mut prefix_set);
         }
     }
 
-    /// Returns a list of paths to the nodes that are located at the provided depth when counting
-    /// from the root node. If there's a leaf at a depth less than the provided depth, it will be
-    /// included in the result.
-    fn get_nodes_at_depth(&self, depth: usize) -> Vec<Nibbles> {
+    /// Returns a list of paths to the nodes that were changed according to the prefix set and are
+    /// located at the provided depth when counting from the root node. If there's a leaf at a
+    /// depth less than the provided depth, it will be included in the result.
+    fn get_changed_nodes_at_depth(&self, prefix_set: &mut PrefixSet, depth: usize) -> Vec<Nibbles> {
         let mut paths = Vec::from([(Nibbles::default(), 0)]);
         let mut targets = Vec::new();
 
         while let Some((mut path, level)) = paths.pop() {
             match self.nodes.get(&path).unwrap() {
                 SparseNode::Empty | SparseNode::Hash(_) => {}
-                SparseNode::Leaf { .. } => {
+                SparseNode::Leaf { hash, .. } => {
+                    if hash.is_some() && !prefix_set.contains(&path) {
+                        continue
+                    }
+
                     targets.push(path);
                 }
-                SparseNode::Extension { key, .. } => {
+                SparseNode::Extension { key, hash } => {
+                    if hash.is_some() && !prefix_set.contains(&path) {
+                        continue
+                    }
+
                     if level >= depth {
                         targets.push(path);
                     } else {
@@ -597,7 +606,11 @@ impl RevealedSparseTrie {
                         paths.push((path, level + 1));
                     }
                 }
-                SparseNode::Branch { state_mask, .. } => {
+                SparseNode::Branch { state_mask, hash } => {
+                    if hash.is_some() && !prefix_set.contains(&path) {
+                        continue
+                    }
+
                     if level >= depth {
                         targets.push(path);
                     } else {
@@ -1540,7 +1553,7 @@ mod tests {
     }
 
     #[test]
-    fn sparse_trie_get_nodes_at_depth() {
+    fn sparse_trie_get_changed_nodes_at_depth() {
         let mut sparse = RevealedSparseTrie::default();
 
         let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
@@ -1574,10 +1587,16 @@ mod tests {
             .unwrap();
         sparse.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x3, 0x2, 0x0]), value).unwrap();
 
-        assert_eq!(sparse.get_nodes_at_depth(0), vec![Nibbles::default()]);
-        assert_eq!(sparse.get_nodes_at_depth(1), vec![Nibbles::from_nibbles_unchecked([0x5])]);
         assert_eq!(
-            sparse.get_nodes_at_depth(2),
+            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 0),
+            vec![Nibbles::default()]
+        );
+        assert_eq!(
+            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 1),
+            vec![Nibbles::from_nibbles_unchecked([0x5])]
+        );
+        assert_eq!(
+            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 2),
             vec![
                 Nibbles::from_nibbles_unchecked([0x5, 0x0]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x2]),
@@ -1585,7 +1604,7 @@ mod tests {
             ]
         );
         assert_eq!(
-            sparse.get_nodes_at_depth(3),
+            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 3),
             vec![
                 Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x2]),
@@ -1594,7 +1613,7 @@ mod tests {
             ]
         );
         assert_eq!(
-            sparse.get_nodes_at_depth(4),
+            sparse.get_changed_nodes_at_depth(&mut PrefixSet::default(), 4),
             vec![
                 Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x1]),
                 Nibbles::from_nibbles_unchecked([0x5, 0x0, 0x2, 0x3, 0x3]),