feat(trie): add update_leaves method to SparseTrieExt (#21525)

Co-authored-by: Amp <amp@ampcode.com>
Co-authored-by: Georgios Konstantopoulos <me@gakonst.com>

crates/trie/sparse/src/provider.rs

@@ -64,6 +64,21 @@ impl TrieNodeProvider for DefaultTrieNodeProvider {
     }
 }
 
+/// A provider that never reveals nodes from the database.
+///
+/// This is used by `update_leaves` to attempt trie operations without
+/// performing any database lookups. When the trie encounters a blinded node
+/// that would normally trigger a reveal, this provider returns `None`,
+/// causing the operation to fail with a `BlindedNode` error.
+#[derive(PartialEq, Eq, Clone, Copy, Default, Debug)]
+pub struct NoRevealProvider;
+
+impl TrieNodeProvider for NoRevealProvider {
+    fn trie_node(&self, _path: &Nibbles) -> Result<Option<RevealedNode>, SparseTrieError> {
+        Ok(None)
+    }
+}
+
 /// Right pad the path with 0s and return as [`B256`].
 #[inline]
 pub fn pad_path_to_key(path: &Nibbles) -> B256 {

crates/trie/sparse/src/traits.rs

@@ -4,7 +4,7 @@ use core::fmt::Debug;
 
 use alloc::{borrow::Cow, vec, vec::Vec};
 use alloy_primitives::{
-    map::{HashMap, HashSet},
+    map::{B256Map, HashMap, HashSet},
     B256,
 };
 use alloy_trie::BranchNodeCompact;
@@ -13,6 +13,17 @@ use reth_trie_common::{BranchNodeMasks, Nibbles, ProofTrieNode, TrieNode};
 
 use crate::provider::TrieNodeProvider;
 
+/// Describes an update to a leaf in the sparse trie.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum LeafUpdate {
+    /// The leaf value has been changed to the given RLP-encoded value.
+    /// Empty Vec indicates the leaf has been removed.
+    Changed(Vec<u8>),
+    /// The leaf value may have changed, but the new value is not yet known.
+    /// Used for optimistic prewarming when the actual value is unavailable.
+    Touched,
+}
+
 /// Trait defining common operations for revealed sparse trie implementations.
 ///
 /// This trait abstracts over different sparse trie implementations (serial vs parallel)
@@ -260,6 +271,26 @@ pub trait SparseTrieExt: SparseTrie {
     ///
     /// The number of nodes converted to hash stubs.
     fn prune(&mut self, max_depth: usize) -> usize;
+
+    /// Applies leaf updates to the sparse trie.
+    ///
+    /// When a [`LeafUpdate::Changed`] is successfully applied, it is removed from the
+    /// given [`B256Map`]. If it could not be applied due to blinded nodes, it remains
+    /// in the map and the callback is invoked with the required proof target.
+    ///
+    /// Once that proof is calculated and revealed via [`SparseTrie::reveal_nodes`], the same
+    /// `updates` map can be reused to retry the update.
+    ///
+    /// Proof targets are deduplicated by `(full_path, min_len)` across all calls to this method.
+    /// The callback will only be invoked once per unique target, even across retry loops.
+    /// A deeper blinded node (higher `min_len`) for the same path is considered a new target.
+    ///
+    /// [`LeafUpdate::Touched`] behaves identically except it does not modify the leaf value.
+    fn update_leaves(
+        &mut self,
+        updates: &mut B256Map<LeafUpdate>,
+        proof_required_fn: impl FnMut(Nibbles, u8),
+    ) -> SparseTrieResult<()>;
 }
 
 /// Tracks modifications to the sparse trie structure.

crates/trie/sparse/src/trie.rs

@@ -1,6 +1,7 @@
 use crate::{
     provider::{RevealedNode, TrieNodeProvider},
-    LeafLookup, LeafLookupError, SparseTrie as SparseTrieTrait, SparseTrieUpdates,
+    LeafLookup, LeafLookupError, LeafUpdate, SparseTrie as SparseTrieTrait, SparseTrieExt,
+    SparseTrieUpdates,
 };
 use alloc::{
     borrow::Cow,
@@ -12,7 +13,7 @@ use alloc::{
 };
 use alloy_primitives::{
     hex, keccak256,
-    map::{Entry, HashMap, HashSet},
+    map::{B256Map, Entry, HashMap, HashSet},
     B256,
 };
 use alloy_rlp::Decodable;
@@ -287,6 +288,36 @@ impl<T: SparseTrieTrait> RevealableSparseTrie<T> {
     }
 }
 
+impl<T: SparseTrieExt + Default> RevealableSparseTrie<T> {
+    /// Applies batch leaf updates to the sparse trie.
+    ///
+    /// For blind tries, all updates are kept in the map and proof targets are emitted
+    /// for every key (with `min_len = 0` since nothing is revealed).
+    ///
+    /// For revealed tries, delegates to the inner implementation which will:
+    /// - Apply updates where possible
+    /// - Keep blocked updates in the map
+    /// - Emit proof targets for blinded paths
+    pub fn update_leaves(
+        &mut self,
+        updates: &mut B256Map<LeafUpdate>,
+        mut proof_required_fn: impl FnMut(Nibbles, u8),
+    ) -> SparseTrieResult<()> {
+        match self {
+            Self::Blind(_) => {
+                // Nothing is revealed - emit proof targets for all keys with min_len = 0
+                for key in updates.keys() {
+                    let full_path = Nibbles::unpack(*key);
+                    proof_required_fn(full_path, 0);
+                }
+                // All updates remain in the map for retry after proofs are fetched
+                Ok(())
+            }
+            Self::Revealed(trie) => trie.update_leaves(updates, proof_required_fn),
+        }
+    }
+}
+
 /// The representation of revealed sparse trie.
 ///
 /// The revealed sparse trie contains the actual trie structure with nodes, values, and