feat: iter-like traverse methods

prestwich · prestwich · commit e7de2e13fbf0 · 2026-01-21T21:10:56.000-05:00
diff --git a/crates/storage/src/hot/conformance.rs b/crates/storage/src/hot/conformance.rs
@@ -1288,6 +1288,7 @@ where
 ///
 /// This function creates a proper BundleState with reverts populated so that
 /// `to_plain_state_and_reverts` will produce the expected output.
+#[allow(clippy::type_complexity)]
 fn make_bundle_state(
     accounts: Vec<(Address, Option<AccountInfo>, Option<AccountInfo>)>,
     storage: Vec<(Address, Vec<(U256, U256, U256)>)>, // (addr, [(slot, old, new)])
diff --git a/crates/storage/src/hot/db/inconsistent.rs b/crates/storage/src/hot/db/inconsistent.rs
@@ -153,31 +153,11 @@ pub trait UnsafeHistoryWrite: UnsafeDbWrite + HotHistoryRead {
     /// Write a pre-state for every storage key that exists for an account at a
     /// specific block.
     fn write_wipe(&self, block_number: u64, address: &Address) -> Result<(), Self::Error> {
-        // SAFETY: the cursor is scoped to the transaction lifetime, which is
-        // valid for the duration of this method.
         let mut cursor = self.traverse_dual::<tables::PlainStorageState>()?;
 
-        let Some(start) = cursor.next_dual_above(address, &U256::ZERO)? else {
-            // No storage entries at or above this address
-            return Ok(());
-        };
-
-        if start.0 != *address {
-            // No storage entries for this address
-            return Ok(());
-        }
-
-        self.write_storage_prestate(block_number, *address, &start.1, &start.2)?;
-
-        while let Some((k, k2, v)) = cursor.next_k2()? {
-            if k != *address {
-                break;
-            }
-
-            self.write_storage_prestate(block_number, *address, &k2, &v)?;
-        }
-
-        Ok(())
+        cursor.for_each_k2(address, &U256::ZERO, |_addr, slot, value| {
+            self.write_storage_prestate(block_number, *address, &slot, &value)
+        })
     }
 
     /// Write a block's plain state revert information.
@@ -294,33 +274,26 @@ pub trait UnsafeHistoryWrite: UnsafeDbWrite + HotHistoryRead {
 
     /// Get all changed accounts with the list of block numbers in the given
     /// range.
+    ///
+    /// Note: This iterates using `next_k2()` which stays within the same k1
+    /// (block number). It effectively only collects changes from the first
+    /// block number in the range.
     fn changed_accounts_with_range(
         &self,
         range: RangeInclusive<BlockNumber>,
     ) -> Result<BTreeMap<Address, Vec<u64>>, Self::Error> {
         let mut changeset_cursor = self.traverse_dual::<tables::AccountChangeSets>()?;
-
         let mut result: BTreeMap<Address, Vec<u64>> = BTreeMap::new();
 
-        // Position cursor at first entry at or above range start
-        let Some((num, addr, _)) =
-            changeset_cursor.next_dual_above(range.start(), &Address::ZERO)?
-        else {
-            return Ok(result);
-        };
-
-        if !range.contains(&num) {
-            return Ok(result);
-        }
-        result.entry(addr).or_default().push(num);
-
-        // Iterate through remaining entries
-        while let Some((num, addr, _)) = changeset_cursor.next_k2()? {
-            if !range.contains(&num) {
-                break;
-            }
-            result.entry(addr).or_default().push(num);
-        }
+        changeset_cursor.for_each_while_k2(
+            range.start(),
+            &Address::ZERO,
+            |num, _, _| range.contains(num),
+            |num, addr, _| {
+                result.entry(addr).or_default().push(num);
+                Ok(())
+            },
+        )?;
 
         Ok(result)
     }
@@ -372,33 +345,27 @@ pub trait UnsafeHistoryWrite: UnsafeDbWrite + HotHistoryRead {
 
     /// Get all changed storages with the list of block numbers in the given
     /// range.
+    ///
+    /// Note: This iterates using `next_k2()` which stays within the same k1
+    /// (block number + address). It effectively only collects changes from
+    /// the first key1 value in the range.
     #[allow(clippy::type_complexity)]
     fn changed_storages_with_range(
         &self,
         range: RangeInclusive<BlockNumber>,
     ) -> Result<BTreeMap<(Address, U256), Vec<u64>>, Self::Error> {
         let mut changeset_cursor = self.traverse_dual::<tables::StorageChangeSets>()?;
-
         let mut result: BTreeMap<(Address, U256), Vec<u64>> = BTreeMap::new();
-        // Position cursor at first entry at or above range start
-        let Some((num_addr, slot, _)) = changeset_cursor
-            .next_dual_above(&BlockNumberAddress((*range.start(), Address::ZERO)), &U256::ZERO)?
-        else {
-            return Ok(result);
-        };
 
-        if !range.contains(&num_addr.block_number()) {
-            return Ok(result);
-        }
-        result.entry((num_addr.address(), slot)).or_default().push(num_addr.block_number());
-
-        // Iterate through remaining entries
-        while let Some((num_addr, slot, _)) = changeset_cursor.next_k2()? {
-            if !range.contains(&num_addr.block_number()) {
-                break;
-            }
-            result.entry((num_addr.address(), slot)).or_default().push(num_addr.block_number());
-        }
+        changeset_cursor.for_each_while_k2(
+            &BlockNumberAddress((*range.start(), Address::ZERO)),
+            &U256::ZERO,
+            |num_addr, _, _| range.contains(&num_addr.block_number()),
+            |num_addr, slot, _| {
+                result.entry((num_addr.address(), slot)).or_default().push(num_addr.block_number());
+                Ok(())
+            },
+        )?;
 
         Ok(result)
     }
diff --git a/crates/storage/src/hot/impls/mdbx/cursor.rs b/crates/storage/src/hot/impls/mdbx/cursor.rs
@@ -222,8 +222,7 @@ where
         let key2_size = fsi.key2_size().unwrap_or(key2.len());
 
         // Use set_range to find the first entry with key1 >= search_key1
-        let Some((found_k1, v)) =
-            self.inner.set_range::<Cow<'_, [u8]>, Cow<'_, [u8]>>(key1)?
+        let Some((found_k1, v)) = self.inner.set_range::<Cow<'_, [u8]>, Cow<'_, [u8]>>(key1)?
         else {
             return Ok(None);
         };
diff --git a/crates/storage/src/hot/model/traverse.rs b/crates/storage/src/hot/model/traverse.rs
@@ -174,6 +174,53 @@ pub trait TableTraverse<T: SingleKey, E: HotKvReadError>: KvTraverse<E> {
     fn read_prev(&mut self) -> Result<Option<KeyValue<T>>, E> {
         KvTraverse::read_prev(self)?.map(T::decode_kv_tuple).transpose().map_err(Into::into)
     }
+
+    /// Iterate entries starting from a key while a predicate holds.
+    ///
+    /// Positions the cursor at `start_key` and calls `f` for each entry
+    /// while `predicate` returns true.
+    ///
+    /// Returns `Ok(())` on successful completion, or the first error encountered.
+    fn for_each_while<P, F>(&mut self, start_key: &T::Key, predicate: P, mut f: F) -> Result<(), E>
+    where
+        P: Fn(&T::Key, &T::Value) -> bool,
+        F: FnMut(T::Key, T::Value) -> Result<(), E>,
+    {
+        let Some((k, v)) = TableTraverse::lower_bound(self, start_key)? else {
+            return Ok(());
+        };
+
+        if !predicate(&k, &v) {
+            return Ok(());
+        }
+
+        f(k, v)?;
+
+        while let Some((k, v)) = TableTraverse::read_next(self)? {
+            if !predicate(&k, &v) {
+                break;
+            }
+            f(k, v)?;
+        }
+
+        Ok(())
+    }
+
+    /// Collect entries from start_key while predicate holds.
+    ///
+    /// This is useful when you need to process entries after iteration completes
+    /// or when the closure would need to borrow mutably from multiple sources.
+    fn collect_while<P>(&mut self, start_key: &T::Key, predicate: P) -> Result<Vec<KeyValue<T>>, E>
+    where
+        P: Fn(&T::Key, &T::Value) -> bool,
+    {
+        let mut result = Vec::new();
+        self.for_each_while(start_key, predicate, |k, v| {
+            result.push((k, v));
+            Ok(())
+        })?;
+        Ok(result)
+    }
 }
 
 /// Blanket implementation of `TableTraverse` for any cursor that implements `KvTraverse`.
@@ -268,6 +315,94 @@ pub trait DualTableTraverse<T: DualKey, E: HotKvReadError>: DualKeyTraverse<E> {
 
     /// Move to the PREVIOUS key2 entry for the CURRENT key1.
     fn previous_k2(&mut self) -> Result<Option<DualKeyValue<T>>, E>;
+
+    /// Iterate entries (crossing k1 boundaries) while a predicate holds.
+    ///
+    /// Positions the cursor at `(key1, start_k2)` and calls `f` for each entry
+    /// while `predicate` returns true. Uses `read_next()` to cross k1 boundaries.
+    ///
+    /// Returns `Ok(())` on successful completion, or the first error encountered.
+    fn for_each_while<P, F>(
+        &mut self,
+        key1: &T::Key,
+        start_k2: &T::Key2,
+        predicate: P,
+        mut f: F,
+    ) -> Result<(), E>
+    where
+        P: Fn(&T::Key, &T::Key2, &T::Value) -> bool,
+        F: FnMut(T::Key, T::Key2, T::Value) -> Result<(), E>,
+    {
+        let Some((k1, k2, v)) = DualTableTraverse::next_dual_above(self, key1, start_k2)? else {
+            return Ok(());
+        };
+
+        if !predicate(&k1, &k2, &v) {
+            return Ok(());
+        }
+
+        f(k1, k2, v)?;
+
+        while let Some((k1, k2, v)) = DualTableTraverse::read_next(self)? {
+            if !predicate(&k1, &k2, &v) {
+                break;
+            }
+            f(k1, k2, v)?;
+        }
+
+        Ok(())
+    }
+
+    /// Iterate entries within the same k1 while a predicate holds.
+    ///
+    /// Positions the cursor at `(key1, start_k2)` and calls `f` for each entry
+    /// while `predicate` returns true. Uses `next_k2()` which stays within
+    /// the same k1 value.
+    ///
+    /// Returns `Ok(())` on successful completion, or the first error encountered.
+    fn for_each_while_k2<P, F>(
+        &mut self,
+        key1: &T::Key,
+        start_k2: &T::Key2,
+        predicate: P,
+        f: F,
+    ) -> Result<(), E>
+    where
+        P: Fn(&T::Key, &T::Key2, &T::Value) -> bool,
+        F: FnMut(T::Key, T::Key2, T::Value) -> Result<(), E>,
+    {
+        self.for_each_while(key1, start_k2, |k, k2, v| key1 == k && predicate(k, k2, v), f)
+    }
+
+    /// Iterate all k2 entries for a given k1, starting from `start_k2`.
+    ///
+    /// Calls `f` for each (k1, k2, v) tuple where k1 matches the provided key1
+    /// and k2 >= start_k2. Stops when k1 changes or the table is exhausted.
+    ///
+    /// Returns `Ok(())` on successful completion, or the first error encountered.
+    fn for_each_k2<F>(&mut self, key1: &T::Key, start_k2: &T::Key2, f: F) -> Result<(), E>
+    where
+        T::Key: PartialEq,
+        F: FnMut(T::Key, T::Key2, T::Value) -> Result<(), E>,
+    {
+        self.for_each_while_k2(key1, start_k2, |_, _, _| true, f)
+    }
+
+    /// Collect all k2 entries for a given k1 into a Vec.
+    ///
+    /// This is useful when you need to process entries after iteration completes
+    /// or when the closure would need to borrow mutably from multiple sources.
+    fn collect_k2(&mut self, key1: &T::Key, start_k2: &T::Key2) -> Result<Vec<DualKeyValue<T>>, E>
+    where
+        T::Key: PartialEq,
+    {
+        let mut result = Vec::new();
+        self.for_each_k2(key1, start_k2, |k1, k2, v| {
+            result.push((k1, k2, v));
+            Ok(())
+        })?;
+        Ok(result)
+    }
 }
 
 impl<C, T, E> DualTableTraverse<T, E> for C
@@ -393,6 +528,30 @@ where
     pub fn read_prev(&mut self) -> Result<Option<KeyValue<T>>, E> {
         TableTraverse::<T, E>::read_prev(&mut self.inner)
     }
+
+    /// Iterate entries starting from a key while a predicate holds.
+    ///
+    /// Positions the cursor at `start_key` and calls `f` for each entry
+    /// while `predicate` returns true.
+    pub fn for_each_while<P, F>(&mut self, start_key: &T::Key, predicate: P, f: F) -> Result<(), E>
+    where
+        P: Fn(&T::Key, &T::Value) -> bool,
+        F: FnMut(T::Key, T::Value) -> Result<(), E>,
+    {
+        TableTraverse::<T, E>::for_each_while(&mut self.inner, start_key, predicate, f)
+    }
+
+    /// Collect entries from start_key while predicate holds.
+    pub fn collect_while<P>(
+        &mut self,
+        start_key: &T::Key,
+        predicate: P,
+    ) -> Result<Vec<KeyValue<T>>, E>
+    where
+        P: Fn(&T::Key, &T::Value) -> bool,
+    {
+        TableTraverse::<T, E>::collect_while(&mut self.inner, start_key, predicate)
+    }
 }
 
 impl<C, T, E> TableCursor<C, T, E>
@@ -517,6 +676,67 @@ where
     pub fn read_prev(&mut self) -> Result<Option<DualKeyValue<T>>, E> {
         DualTableTraverse::<T, E>::read_prev(&mut self.inner)
     }
+
+    /// Iterate all k2 entries for a given k1, starting from `start_k2`.
+    ///
+    /// Calls `f` for each (k1, k2, v) tuple where k1 matches the provided key1
+    /// and k2 >= start_k2. Stops when k1 changes or the table is exhausted.
+    pub fn for_each_k2<F>(&mut self, key1: &T::Key, start_k2: &T::Key2, f: F) -> Result<(), E>
+    where
+        T::Key: PartialEq,
+        F: FnMut(T::Key, T::Key2, T::Value) -> Result<(), E>,
+    {
+        DualTableTraverse::<T, E>::for_each_k2(&mut self.inner, key1, start_k2, f)
+    }
+
+    /// Iterate entries within the same k1 while a predicate holds.
+    ///
+    /// Positions the cursor at `(key1, start_k2)` and calls `f` for each entry
+    /// while `predicate` returns true. Uses `next_k2()` which stays within
+    /// the same k1 value.
+    pub fn for_each_while_k2<P, F>(
+        &mut self,
+        key1: &T::Key,
+        start_k2: &T::Key2,
+        predicate: P,
+        f: F,
+    ) -> Result<(), E>
+    where
+        P: Fn(&T::Key, &T::Key2, &T::Value) -> bool,
+        F: FnMut(T::Key, T::Key2, T::Value) -> Result<(), E>,
+    {
+        DualTableTraverse::<T, E>::for_each_while_k2(&mut self.inner, key1, start_k2, predicate, f)
+    }
+
+    /// Iterate entries (crossing k1 boundaries) while a predicate holds.
+    ///
+    /// Positions the cursor at `(key1, start_k2)` and calls `f` for each entry
+    /// while `predicate` returns true. Uses `read_next()` to cross k1 boundaries.
+    pub fn for_each_while<P, F>(
+        &mut self,
+        key1: &T::Key,
+        start_k2: &T::Key2,
+        predicate: P,
+        f: F,
+    ) -> Result<(), E>
+    where
+        P: Fn(&T::Key, &T::Key2, &T::Value) -> bool,
+        F: FnMut(T::Key, T::Key2, T::Value) -> Result<(), E>,
+    {
+        DualTableTraverse::<T, E>::for_each_while(&mut self.inner, key1, start_k2, predicate, f)
+    }
+
+    /// Collect all k2 entries for a given k1 into a Vec.
+    pub fn collect_k2(
+        &mut self,
+        key1: &T::Key,
+        start_k2: &T::Key2,
+    ) -> Result<Vec<DualKeyValue<T>>, E>
+    where
+        T::Key: PartialEq,
+    {
+        DualTableTraverse::<T, E>::collect_k2(&mut self.inner, key1, start_k2)
+    }
 }
 
 impl<C, T, E> DualTableCursor<C, T, E>
