feat(trie): add clear method to PrefixSetMut, RevealedSparseTrie, SparseTrieUpdates (#16179)

Co-authored-by: Alexey Shekhirin <[email protected]>

Author: Rjected

crates/trie/common/src/prefix_set.rs

@@ -136,6 +136,12 @@ impl PrefixSetMut {
         self.keys.is_empty()
     }
 
+    /// Clears the inner vec for reuse, setting `all` to `false`.
+    pub fn clear(&mut self) {
+        self.all = false;
+        self.keys.clear();
+    }
+
     /// Returns a `PrefixSet` with the same elements as this set.
     ///
     /// If not yet sorted, the elements will be sorted and deduplicated.

crates/trie/sparse/src/trie.rs

@@ -1317,6 +1317,22 @@ pub enum LeafLookup {
 }
 
 impl<P: BlindedProvider> RevealedSparseTrie<P> {
+    /// This clears all data structures in the sparse trie, keeping the backing data structures
+    /// allocated.
+    ///
+    /// This is useful for reusing the trie without needing to reallocate memory.
+    pub fn clear(&mut self) {
+        self.nodes.clear();
+        self.branch_node_tree_masks.clear();
+        self.branch_node_hash_masks.clear();
+        self.values.clear();
+        self.prefix_set.clear();
+        if let Some(updates) = self.updates.as_mut() {
+            updates.clear()
+        }
+        self.rlp_buf.clear();
+    }
+
     /// Attempts to find a leaf node at the specified path.
     ///
     /// This method traverses the trie from the root down to the given path, checking
@@ -2019,6 +2035,15 @@ impl SparseTrieUpdates {
     pub fn wiped() -> Self {
         Self { wiped: true, ..Default::default() }
     }
+
+    /// Clears the updates, but keeps the backing data structures allocated.
+    ///
+    /// Sets `wiped` to `false`.
+    pub fn clear(&mut self) {
+        self.updated_nodes.clear();
+        self.removed_nodes.clear();
+        self.wiped = false;
+    }
 }
 
 #[cfg(test)]
@@ -3616,6 +3641,35 @@ mod tests {
         assert_eq!(sparse.root(), EMPTY_ROOT_HASH);
     }
 
+    #[test]
+    fn sparse_trie_clear() {
+        // tests that if we fill a sparse trie with some nodes and then clear it, it has the same
+        // contents as an empty sparse trie
+        let mut sparse = RevealedSparseTrie::default();
+        let value = alloy_rlp::encode_fixed_size(&U256::ZERO).to_vec();
+        sparse
+            .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x1]), value.clone())
+            .unwrap();
+        sparse
+            .update_leaf(Nibbles::from_nibbles([0x5, 0x0, 0x2, 0x3, 0x3]), value.clone())
+            .unwrap();
+        sparse
+            .update_leaf(Nibbles::from_nibbles([0x5, 0x2, 0x0, 0x1, 0x3]), value.clone())
+            .unwrap();
+        sparse.update_leaf(Nibbles::from_nibbles([0x5, 0x3, 0x1, 0x0, 0x2]), value).unwrap();
+
+        sparse.clear();
+
+        // we have to update the root hash to be an empty one, because the `Default` impl of
+        // `RevealedSparseTrie` sets the root hash to `EMPTY_ROOT_HASH` in the constructor.
+        //
+        // The default impl is only used in tests.
+        sparse.nodes.insert(Nibbles::default(), SparseNode::Empty);
+
+        let empty_trie = RevealedSparseTrie::default();
+        assert_eq!(empty_trie, sparse);
+    }
+
     #[test]
     fn sparse_trie_display() {
         let mut sparse = RevealedSparseTrie::default();