Skip to content

Commit

Permalink
feat(storage): add db table iterator
Browse files Browse the repository at this point in the history
  • Loading branch information
Yael-Starkware committed Oct 2, 2023
1 parent e5a8d1b commit cb3d015
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 1 deletion.
37 changes: 36 additions & 1 deletion crates/papyrus_storage/src/db/db_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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::<DbResult<Vec<_>>>().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;
}
}
27 changes: 27 additions & 0 deletions crates/papyrus_storage/src/db/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<K>,
_value_type: PhantomData<V>,
}

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<Self::Item> {
let prev_cursor_res = self.cursor.next().transpose()?;
Some(prev_cursor_res)
}
}

#[doc(hidden)]
pub struct RO {}

Expand Down

0 comments on commit cb3d015

Please sign in to comment.