perf(trie): optimize in memory cursor (#14600)

crates/trie/trie/src/forward_cursor.rs

@@ -1,23 +1,36 @@
 /// The implementation of forward-only in memory cursor over the entries.
+///
 /// The cursor operates under the assumption that the supplied collection is pre-sorted.
 #[derive(Debug)]
 pub struct ForwardInMemoryCursor<'a, K, V> {
     /// The reference to the pre-sorted collection of entries.
-    entries: &'a Vec<(K, V)>,
+    entries: std::slice::Iter<'a, (K, V)>,
-    /// The index where cursor is currently positioned.
+    is_empty: bool,
-    index: usize,
 }
 
 impl<'a, K, V> ForwardInMemoryCursor<'a, K, V> {
     /// Create new forward cursor positioned at the beginning of the collection.
+    ///
     /// The cursor expects all of the entries have been sorted in advance.
-    pub const fn new(entries: &'a Vec<(K, V)>) -> Self {
+    #[inline]
-        Self { entries, index: 0 }
+    pub fn new(entries: &'a [(K, V)]) -> Self {
+        Self { entries: entries.iter(), is_empty: entries.is_empty() }
     }
 
     /// Returns `true` if the cursor is empty, regardless of its position.
+    #[inline]
     pub fn is_empty(&self) -> bool {
-        self.entries.is_empty()
+        self.is_empty
+    }
+
+    #[inline]
+    fn peek(&self) -> Option<&(K, V)> {
+        self.entries.clone().next()
+    }
+
+    #[inline]
+    fn next(&mut self) -> Option<&(K, V)> {
+        self.entries.next()
     }
 }
 
@@ -26,26 +39,58 @@ where
     K: PartialOrd + Clone,
     V: Clone,
 {
-    /// Advances the cursor forward while `comparator` returns `true` or until the collection is
-    /// exhausted. Returns the first entry for which `comparator` returns `false` or `None`.
-    fn advance_while_false(&mut self, comparator: impl Fn(&K) -> bool) -> Option<(K, V)> {
-        let mut entry = self.entries.get(self.index);
-        while entry.is_some_and(|entry| comparator(&entry.0)) {
-            self.index += 1;
-            entry = self.entries.get(self.index);
-        }
-        entry.cloned()
-    }
-
     /// Returns the first entry from the current cursor position that's greater or equal to the
     /// provided key. This method advances the cursor forward.
     pub fn seek(&mut self, key: &K) -> Option<(K, V)> {
-        self.advance_while_false(|k| k < key)
+        self.advance_while(|k| k < key)
     }
 
     /// Returns the first entry from the current cursor position that's greater than the provided
     /// key. This method advances the cursor forward.
     pub fn first_after(&mut self, key: &K) -> Option<(K, V)> {
-        self.advance_while_false(|k| k <= key)
+        self.advance_while(|k| k <= key)
+    }
+
+    /// Advances the cursor forward while `predicate` returns `true` or until the collection is
+    /// exhausted.
+    ///
+    /// Returns the first entry for which `predicate` returns `false` or `None`. The cursor will
+    /// point to the returned entry.
+    fn advance_while(&mut self, predicate: impl Fn(&K) -> bool) -> Option<(K, V)> {
+        let mut entry;
+        loop {
+            entry = self.peek();
+            if entry.is_some_and(|(k, _)| predicate(k)) {
+                self.next();
+            } else {
+                break;
+            }
+        }
+        entry.cloned()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_cursor() {
+        let mut cursor = ForwardInMemoryCursor::new(&[(1, ()), (2, ()), (3, ()), (4, ()), (5, ())]);
+
+        assert_eq!(cursor.seek(&0), Some((1, ())));
+        assert_eq!(cursor.peek(), Some(&(1, ())));
+
+        assert_eq!(cursor.seek(&3), Some((3, ())));
+        assert_eq!(cursor.peek(), Some(&(3, ())));
+
+        assert_eq!(cursor.seek(&3), Some((3, ())));
+        assert_eq!(cursor.peek(), Some(&(3, ())));
+
+        assert_eq!(cursor.seek(&4), Some((4, ())));
+        assert_eq!(cursor.peek(), Some(&(4, ())));
+
+        assert_eq!(cursor.seek(&6), None);
+        assert_eq!(cursor.peek(), None);
     }
 }

crates/trie/trie/src/hashed_cursor/post_state.rs

@@ -57,7 +57,7 @@ where
     C: HashedCursor<Value = Account>,
 {
     /// Create new instance of [`HashedPostStateAccountCursor`].
-    pub const fn new(cursor: C, post_state_accounts: &'a HashedAccountsSorted) -> Self {
+    pub fn new(cursor: C, post_state_accounts: &'a HashedAccountsSorted) -> Self {
         let post_state_cursor = ForwardInMemoryCursor::new(&post_state_accounts.accounts);
         let destroyed_accounts = &post_state_accounts.destroyed_accounts;
         Self { cursor, post_state_cursor, destroyed_accounts, last_account: None }

crates/trie/trie/src/trie_cursor/in_memory.rs

@@ -62,7 +62,7 @@ pub struct InMemoryAccountTrieCursor<'a, C> {
 impl<'a, C: TrieCursor> InMemoryAccountTrieCursor<'a, C> {
     /// Create new account trie cursor from underlying cursor and reference to
     /// [`TrieUpdatesSorted`].
-    pub const fn new(cursor: C, trie_updates: &'a TrieUpdatesSorted) -> Self {
+    pub fn new(cursor: C, trie_updates: &'a TrieUpdatesSorted) -> Self {
         let in_memory_cursor = ForwardInMemoryCursor::new(&trie_updates.account_nodes);
         Self {
             cursor,