feat(trie): proof_v2 prefix set support (#22946)

Co-authored-by: Amp <amp@ampcode.com>

.changelog/dry-lakes-sleep.md

@@ -0,0 +1,6 @@
+---
+reth-trie-common: minor
+reth-trie: minor
+---
+
+Added `contains_range` method to `PrefixSet` for checking if any key falls within a half-open range. Added prefix set support to `ProofCalculator` via `with_prefix_set`, enabling stale cached hash invalidation and branch collapse detection when keys are inserted or removed; propagated storage prefix sets through `SyncAccountValueEncoder`.

crates/trie/common/src/prefix_set.rs

@@ -1,6 +1,7 @@
 use crate::Nibbles;
 use alloc::{sync::Arc, vec::Vec};
 use alloy_primitives::map::{B256Map, B256Set};
+use core::ops::Range;
 
 /// Collection of mutable prefix sets.
 #[derive(Clone, Default, Debug)]
@@ -225,6 +226,36 @@ impl PrefixSet {
         false
     }
 
+    /// Returns `true` if any key in the set falls within the given half-open range
+    /// `[start, end)`.
+    ///
+    /// Like [`Self::contains`], this method maintains the internal index for sequential access
+    /// optimization.
+    #[inline]
+    pub fn contains_range(&mut self, range: Range<&Nibbles>) -> bool {
+        if self.all {
+            return true
+        }
+
+        while self.index > 0 && &self.keys[self.index] >= range.end {
+            self.index -= 1;
+        }
+
+        for (idx, key) in self.keys[self.index..].iter().enumerate() {
+            if key >= range.start && key < range.end {
+                self.index += idx;
+                return true
+            }
+
+            if key >= range.end {
+                self.index += idx;
+                return false
+            }
+        }
+
+        false
+    }
+
     /// Returns an iterator over reference to _all_ nibbles regardless of cursor position.
     pub fn iter(&self) -> core::slice::Iter<'_, Nibbles> {
         self.keys.iter()

crates/trie/trie/src/proof_v2/value.rs

@@ -1,9 +1,10 @@
 //! Generic value encoder types for proof calculation with lazy evaluation.
 
 use crate::{
-    hashed_cursor::HashedCursorFactory, proof_v2::ProofCalculator, trie_cursor::TrieCursorFactory,
+    hashed_cursor::HashedCursorFactory, prefix_set::PrefixSet, proof_v2::ProofCalculator,
+    trie_cursor::TrieCursorFactory,
 };
-use alloy_primitives::{B256, U256};
+use alloy_primitives::{map::B256Map, B256, U256};
 use alloy_rlp::Encodable;
 use reth_execution_errors::trie::StateProofError;
 use reth_primitives_traits::Account;
@@ -83,6 +84,8 @@ pub struct SyncAccountValueEncoder<T, H> {
     trie_cursor_factory: Rc<T>,
     /// Factory for creating hashed cursors.
     hashed_cursor_factory: Rc<H>,
+    /// Storage prefix sets keyed by hashed address.
+    storage_prefix_sets: Rc<B256Map<PrefixSet>>,
 }
 
 impl<T, H> SyncAccountValueEncoder<T, H> {
@@ -91,8 +94,17 @@ impl<T, H> SyncAccountValueEncoder<T, H> {
         Self {
             trie_cursor_factory: Rc::new(trie_cursor_factory),
             hashed_cursor_factory: Rc::new(hashed_cursor_factory),
+            storage_prefix_sets: Rc::new(B256Map::default()),
         }
     }
+
+    /// Sets the storage prefix sets. When given, all cached storage trie hashes matching the
+    /// prefix sets will be invalidated during storage root calculation for the corresponding
+    /// accounts.
+    pub fn with_storage_prefix_sets(mut self, storage_prefix_sets: B256Map<PrefixSet>) -> Self {
+        self.storage_prefix_sets = Rc::new(storage_prefix_sets);
+        self
+    }
 }
 
 /// The deferred encoder for an account value with synchronous storage root calculation.
@@ -100,6 +112,7 @@ impl<T, H> SyncAccountValueEncoder<T, H> {
 pub struct SyncAccountDeferredValueEncoder<T, H> {
     trie_cursor_factory: Rc<T>,
     hashed_cursor_factory: Rc<H>,
+    storage_prefix_sets: Rc<B256Map<PrefixSet>>,
     hashed_address: B256,
     account: Account,
 }
@@ -115,6 +128,9 @@ where
             self.hashed_cursor_factory.hashed_storage_cursor(self.hashed_address)?;
 
         let mut storage_proof_calculator = ProofCalculator::new_storage(trie_cursor, hashed_cursor);
+        if let Some(prefix_set) = self.storage_prefix_sets.get(&self.hashed_address) {
+            storage_proof_calculator = storage_proof_calculator.with_prefix_set(prefix_set.clone());
+        }
         let root_node = storage_proof_calculator.storage_root_node(self.hashed_address)?;
         let storage_root = storage_proof_calculator
             .compute_root_hash(&[root_node])?
@@ -143,8 +159,9 @@ where
         // Return a deferred encoder that will synchronously compute the storage root when encode()
         // is called.
         SyncAccountDeferredValueEncoder {
-            trie_cursor_factory: self.trie_cursor_factory.clone(),
-            hashed_cursor_factory: self.hashed_cursor_factory.clone(),
+            trie_cursor_factory: Rc::clone(&self.trie_cursor_factory),
+            hashed_cursor_factory: Rc::clone(&self.hashed_cursor_factory),
+            storage_prefix_sets: Rc::clone(&self.storage_prefix_sets),
             hashed_address,
             account,
         }