diff --git a/crates/papyrus_storage/src/db/db_test.rs b/crates/papyrus_storage/src/db/db_test.rs index f17cbbb239..452c73625d 100644 --- a/crates/papyrus_storage/src/db/db_test.rs +++ b/crates/papyrus_storage/src/db/db_test.rs @@ -2,7 +2,7 @@ use libmdbx::PageSize; use pretty_assertions::assert_eq; use tempfile::TempDir; -use crate::db::{get_page_size, open_env, DbReader, DbWriter}; +use crate::db::{get_page_size, open_env, DbIter, DbReader, DbResult, DbWriter}; use crate::test_utils::get_test_config; fn get_test_env() -> ((DbReader, DbWriter), TempDir) { @@ -115,3 +115,38 @@ fn get_page_size_test() { assert_eq!(get_page_size(1025), PageSize::Set(1024)); assert_eq!(get_page_size(2047), PageSize::Set(1024)); } + +#[test] +fn test_iter() { + // Create an environment and a table. + let ((reader, mut writer), _temp_dir) = get_test_env(); + let table_id = writer.create_table::<[u8; 4], [u8; 4]>("table").unwrap(); + + // Insert some values. + let items = vec![ + (*b"key1", *b"val1"), + (*b"key2", *b"val2"), + (*b"key3", *b"val3"), + (*b"key5", *b"val5"), + ]; + let wtxn = writer.begin_rw_txn().unwrap(); + let table = wtxn.open_table(&table_id).unwrap(); + for (k, v) in &items { + table.insert(&wtxn, k, v).unwrap(); + } + wtxn.commit().unwrap(); + + // Use the iterator. + let txn = reader.begin_ro_txn().unwrap(); + let mut cursor = txn.open_table(&table_id).unwrap().cursor(&txn).unwrap(); + let iter = DbIter::new(&mut cursor); + assert_eq!(items, iter.collect::>>().unwrap()); + + let mut cursor = txn.open_table(&table_id).unwrap().cursor(&txn).unwrap(); + let mut iter = DbIter::new(&mut cursor); + let mut index = 0; + while let Some(Ok((k, v))) = iter.next() { + assert_eq!(items[index], (k, v)); + index += 1; + } +} diff --git a/crates/papyrus_storage/src/db/mod.rs b/crates/papyrus_storage/src/db/mod.rs index 730de2d15e..89a5851fe2 100644 --- a/crates/papyrus_storage/src/db/mod.rs +++ b/crates/papyrus_storage/src/db/mod.rs @@ -380,6 +380,33 @@ impl<'txn, Mode: TransactionKind, K: StorageSerde, V: StorageSerde> DbCursor<'tx } } +/// A wrapper around a cursor that allows iterating over a table. +pub(crate) struct DbIter<'cursor, 'txn, Mode: TransactionKind, K: StorageSerde, V: StorageSerde> { + cursor: &'cursor mut DbCursor<'txn, Mode, K, V>, + _key_type: PhantomData, + _value_type: PhantomData, +} + +impl<'cursor, 'txn, Mode: TransactionKind, K: StorageSerde, V: StorageSerde> + DbIter<'cursor, 'txn, Mode, K, V> +{ + #[allow(dead_code)] + pub(crate) fn new(cursor: &'cursor mut DbCursor<'txn, Mode, K, V>) -> Self { + Self { cursor, _key_type: PhantomData {}, _value_type: PhantomData {} } + } +} + +impl<'cursor, 'txn, Mode: TransactionKind, K: StorageSerde, V: StorageSerde> Iterator + for DbIter<'cursor, 'txn, Mode, K, V> +{ + type Item = DbResult<(K, V)>; + + fn next(&mut self) -> Option { + let prev_cursor_res = self.cursor.next().transpose()?; + Some(prev_cursor_res) + } +} + #[doc(hidden)] pub struct RO {}