feat(trie): add prefetch for blinded nodes to reduce random I/O

Amp-Thread-ID: https://ampcode.com/threads/T-019bfe25-43f3-75ac-98f7-32bf937b69e1
Co-authored-by: Amp <amp@ampcode.com>

crates/trie/sparse/src/state.rs

@@ -224,6 +224,31 @@ where
         self.storage.tries.insert(address, storage_trie);
     }
 
+    /// Pre-reveal blinded nodes for account paths to reduce random I/O.
+    pub fn prefetch_blinded_for_paths<P: TrieNodeProvider>(
+        &mut self,
+        account_paths: impl Iterator<Item = Nibbles>,
+        provider: &P,
+    ) -> SparseStateTrieResult<()>
+    where
+        A: crate::BlindedPathCollector,
+    {
+        let mut blinded = Vec::new();
+        if let Some(trie) = self.state.as_revealed_ref() {
+            for path in account_paths {
+                trie.collect_blinded_on_path(&path, &mut blinded);
+            }
+        }
+        for path in blinded {
+            if let Some(node) = provider.trie_node(&path)? {
+                if let Some(trie) = self.state.as_revealed_mut() {
+                    trie.reveal_node(path, TrieNode::decode(&mut node.node.as_ref())?, None)?;
+                }
+            }
+        }
+        Ok(())
+    }
+
     /// Reveal unknown trie paths from multiproof.
     /// NOTE: This method does not extensively validate the proof.
     pub fn reveal_multiproof(&mut self, multiproof: MultiProof) -> SparseStateTrieResult<()> {

crates/trie/sparse/src/traits.rs

@@ -277,3 +277,9 @@ pub enum LeafLookup {
     /// Leaf does not exist (exclusion proof found).
     NonExistent,
 }
+
+/// Trait for collecting blinded node paths along a target path.
+pub trait BlindedPathCollector {
+    /// Collect paths of blinded (Hash) nodes along the path from root to target.
+    fn collect_blinded_on_path(&self, target_path: &Nibbles, out: &mut Vec<Nibbles>);
+}

crates/trie/sparse/src/trie.rs

@@ -1,6 +1,7 @@
 use crate::{
     provider::{RevealedNode, TrieNodeProvider},
-    LeafLookup, LeafLookupError, SparseTrie as SparseTrieTrait, SparseTrieUpdates,
+    BlindedPathCollector, LeafLookup, LeafLookupError, SparseTrie as SparseTrieTrait,
+    SparseTrieUpdates,
 };
 use alloc::{
     borrow::Cow,
@@ -421,6 +422,12 @@ impl Default for SerialSparseTrie {
     }
 }
 
+impl BlindedPathCollector for SerialSparseTrie {
+    fn collect_blinded_on_path(&self, target_path: &Nibbles, out: &mut Vec<Nibbles>) {
+        self.collect_blinded_on_path(target_path, out)
+    }
+}
+
 impl SparseTrieTrait for SerialSparseTrie {
     fn with_root(
         mut self,
@@ -1133,6 +1140,34 @@ impl SerialSparseTrie {
         &self.nodes
     }
 
+    /// Collect paths of blinded (Hash) nodes along the path from root to target.
+    pub fn collect_blinded_on_path(&self, target_path: &Nibbles, out: &mut Vec<Nibbles>) {
+        let mut current = Nibbles::default();
+        while current.len() < target_path.len() {
+            if let Some(node) = self.nodes.get(&current) {
+                if matches!(node, SparseNode::Hash(_)) {
+                    out.push(current.clone());
+                }
+                match node {
+                    SparseNode::Branch { state_mask, .. } => {
+                        let nibble = target_path.get_unchecked(current.len());
+                        if state_mask.is_bit_set(nibble) {
+                            current.push_unchecked(nibble);
+                        } else {
+                            break;
+                        }
+                    }
+                    SparseNode::Extension { key, .. } => {
+                        current.extend(key);
+                    }
+                    _ => break,
+                }
+            } else {
+                break;
+            }
+        }
+    }
+
     /// Reveals either a node or its hash placeholder based on the provided child data.
     ///
     /// When traversing the trie, we often encounter references to child nodes that