From 9482c551c59443e2e4da8a9046e4ebfb0ce24846 Mon Sep 17 00:00:00 2001 From: TurboFish <87058147+CryptoTurboFish@users.noreply.github.com> Date: Mon, 16 Jan 2023 10:27:22 -0800 Subject: [PATCH] feat: add a reverse db walker (#841) --- bin/reth/src/db/mod.rs | 6 +- crates/storage/db/src/abstraction/cursor.rs | 58 ++++++++++++++++- .../db/src/implementation/mdbx/cursor.rs | 2 +- .../storage/db/src/implementation/mdbx/mod.rs | 62 ++++++++++++++++++- 4 files changed, 118 insertions(+), 10 deletions(-) diff --git a/bin/reth/src/db/mod.rs b/bin/reth/src/db/mod.rs index b79e8ddf53..003c3502f2 100644 --- a/bin/reth/src/db/mod.rs +++ b/bin/reth/src/db/mod.rs @@ -190,11 +190,7 @@ impl<'a, DB: Database> DbTool<'a, DB> { // TODO: Upstream this in the DB trait. let start_walker = cursor.current().transpose(); - let walker = Walker { - cursor: &mut cursor, - start: start_walker, - _tx_phantom: std::marker::PhantomData, - }; + let walker = Walker::new(&mut cursor, start_walker); walker.skip(start).take(len).collect::>() })?; diff --git a/crates/storage/db/src/abstraction/cursor.rs b/crates/storage/db/src/abstraction/cursor.rs index fda0c2c9cf..8a3414c8c4 100644 --- a/crates/storage/db/src/abstraction/cursor.rs +++ b/crates/storage/db/src/abstraction/cursor.rs @@ -101,11 +101,11 @@ pub trait DbDupCursorRW<'tx, T: DupSort> { /// the Cursor lifetime and it wouldn't be possible to use Walker. pub struct Walker<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> { /// Cursor to be used to walk through the table. - pub cursor: &'cursor mut CURSOR, + cursor: &'cursor mut CURSOR, /// `(key, value)` where to start the walk. - pub start: IterPairResult, + start: IterPairResult, /// Phantom data for 'tx. As it is only used for `DbCursorRO`. - pub _tx_phantom: PhantomData<&'tx T>, + _tx_phantom: PhantomData<&'tx T>, } impl<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> std::iter::Iterator @@ -122,6 +122,58 @@ impl<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> std::iter::Iterator } } +impl<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> Walker<'cursor, 'tx, T, CURSOR> { + /// construct Walker + pub fn new(cursor: &'cursor mut CURSOR, start: IterPairResult) -> Self { + Self { cursor, start, _tx_phantom: std::marker::PhantomData } + } + + /// convert current [`Walker`] to [`ReverseWalker`] which iterates reversely + pub fn rev(self) -> ReverseWalker<'cursor, 'tx, T, CURSOR> { + let start = self.cursor.current().transpose(); + ReverseWalker::new(self.cursor, start) + } +} + +/// Provides a reverse iterator to `Cursor` when handling `Table`. +/// Also check [`Walker`] +pub struct ReverseWalker<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> { + /// Cursor to be used to walk through the table. + cursor: &'cursor mut CURSOR, + /// `(key, value)` where to start the walk. + start: IterPairResult, + /// Phantom data for 'tx. As it is only used for `DbCursorRO`. + _tx_phantom: PhantomData<&'tx T>, +} + +impl<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> ReverseWalker<'cursor, 'tx, T, CURSOR> { + /// construct ReverseWalker + pub fn new(cursor: &'cursor mut CURSOR, start: IterPairResult) -> Self { + Self { cursor, start, _tx_phantom: std::marker::PhantomData } + } + + /// convert current [`ReverseWalker`] to [`Walker`] which iterate forwardly + pub fn forward(self) -> Walker<'cursor, 'tx, T, CURSOR> { + let start = self.cursor.current().transpose(); + Walker::new(self.cursor, start) + } +} + +impl<'cursor, 'tx, T: Table, CURSOR: DbCursorRO<'tx, T>> std::iter::Iterator + for ReverseWalker<'cursor, 'tx, T, CURSOR> +{ + type Item = Result<(T::Key, T::Value), Error>; + + fn next(&mut self) -> Option { + let start = self.start.take(); + if start.is_some() { + return start + } + + self.cursor.prev().transpose() + } +} + /// Provides an iterator to `Cursor` when handling a `DupSort` table. /// /// Reason why we have two lifetimes is to distinguish between `'cursor` lifetime diff --git a/crates/storage/db/src/implementation/mdbx/cursor.rs b/crates/storage/db/src/implementation/mdbx/cursor.rs index 586cba7f86..61709bdcf3 100644 --- a/crates/storage/db/src/implementation/mdbx/cursor.rs +++ b/crates/storage/db/src/implementation/mdbx/cursor.rs @@ -79,7 +79,7 @@ impl<'tx, K: TransactionKind, T: Table> DbCursorRO<'tx, T> for Cursor<'tx, K, T> .map_err(|e| Error::Read(e.into()))? .map(decoder::); - Ok(Walker::<'cursor, 'tx, T, Self> { cursor: self, start, _tx_phantom: PhantomData {} }) + Ok(Walker::new(self, start)) } } diff --git a/crates/storage/db/src/implementation/mdbx/mod.rs b/crates/storage/db/src/implementation/mdbx/mod.rs index f7292a711a..c5c3e3da14 100644 --- a/crates/storage/db/src/implementation/mdbx/mod.rs +++ b/crates/storage/db/src/implementation/mdbx/mod.rs @@ -144,7 +144,7 @@ pub mod test_utils { mod tests { use super::{test_utils, Env, EnvKind}; use crate::{ - cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO}, + cursor::{DbCursorRO, DbCursorRW, DbDupCursorRO, ReverseWalker, Walker}, database::Database, models::ShardedKey, tables::{AccountHistory, CanonicalHeaders, Headers, PlainAccountState, PlainStorageState}, @@ -213,6 +213,66 @@ mod tests { assert_eq!(first.1, value, "First next should be put value"); } + #[test] + fn db_walker() { + let db: Arc> = test_utils::create_test_db(EnvKind::RW); + + // PUT (0, 0), (1, 0), (3, 0) + let tx = db.tx_mut().expect(ERROR_INIT_TX); + vec![0, 1, 3] + .into_iter() + .try_for_each(|key| tx.put::(key, H256::zero())) + .expect(ERROR_PUT); + tx.commit().expect(ERROR_COMMIT); + + let tx = db.tx().expect(ERROR_INIT_TX); + let mut cursor = tx.cursor_read::().unwrap(); + + let mut walker = Walker::new(&mut cursor, None); + + assert_eq!(walker.next(), Some(Ok((0, H256::zero())))); + assert_eq!(walker.next(), Some(Ok((1, H256::zero())))); + assert_eq!(walker.next(), Some(Ok((3, H256::zero())))); + assert_eq!(walker.next(), None); + + // transform to ReverseWalker + let mut reverse_walker = walker.rev(); + assert_eq!(reverse_walker.next(), Some(Ok((3, H256::zero())))); + assert_eq!(reverse_walker.next(), Some(Ok((1, H256::zero())))); + assert_eq!(reverse_walker.next(), Some(Ok((0, H256::zero())))); + assert_eq!(reverse_walker.next(), None); + } + + #[test] + fn db_reverse_walker() { + let db: Arc> = test_utils::create_test_db(EnvKind::RW); + + // PUT (0, 0), (1, 0), (3, 0) + let tx = db.tx_mut().expect(ERROR_INIT_TX); + vec![0, 1, 3] + .into_iter() + .try_for_each(|key| tx.put::(key, H256::zero())) + .expect(ERROR_PUT); + tx.commit().expect(ERROR_COMMIT); + + let tx = db.tx().expect(ERROR_INIT_TX); + let mut cursor = tx.cursor_read::().unwrap(); + + let mut reverse_walker = ReverseWalker::new(&mut cursor, None); + + assert_eq!(reverse_walker.next(), Some(Ok((3, H256::zero())))); + assert_eq!(reverse_walker.next(), Some(Ok((1, H256::zero())))); + assert_eq!(reverse_walker.next(), Some(Ok((0, H256::zero())))); + assert_eq!(reverse_walker.next(), None); + + // transform to Walker + let mut walker = reverse_walker.forward(); + assert_eq!(walker.next(), Some(Ok((0, H256::zero())))); + assert_eq!(walker.next(), Some(Ok((1, H256::zero())))); + assert_eq!(walker.next(), Some(Ok((3, H256::zero())))); + assert_eq!(walker.next(), None); + } + #[test] fn db_cursor_seek_exact_or_previous_key() { let db: Arc> = test_utils::create_test_db(EnvKind::RW);