diff --git a/crates/trie/trie/src/forward_cursor.rs b/crates/trie/trie/src/forward_cursor.rs index 745fc351b..6f7cdbb42 100644 --- a/crates/trie/trie/src/forward_cursor.rs +++ b/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)>, - /// The index where cursor is currently positioned. - index: usize, + entries: std::slice::Iter<'a, (K, V)>, + is_empty: bool, } 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 { - Self { entries, index: 0 } + #[inline] + 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); } } diff --git a/crates/trie/trie/src/hashed_cursor/post_state.rs b/crates/trie/trie/src/hashed_cursor/post_state.rs index 31e760e8c..b6c8994e1 100644 --- a/crates/trie/trie/src/hashed_cursor/post_state.rs +++ b/crates/trie/trie/src/hashed_cursor/post_state.rs @@ -57,7 +57,7 @@ where C: HashedCursor, { /// 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 } diff --git a/crates/trie/trie/src/trie_cursor/in_memory.rs b/crates/trie/trie/src/trie_cursor/in_memory.rs index 0e394fcf7..17c5d3b4a 100644 --- a/crates/trie/trie/src/trie_cursor/in_memory.rs +++ b/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,