Merge branch 'origin/ray/iter-optional-prefix' (#1478)

* origin/ray/iter-optional-prefix: storage: add optional prefix iterator methods
anoma · Jun 14, 2023 · 0f04641 · 0f04641
2 parents 20d8647 + 7aaf3ce
commit 0f04641
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 7 deletions.
diff --git a/apps/src/lib/node/ledger/storage/rocksdb.rs b/apps/src/lib/node/ledger/storage/rocksdb.rs
@@ -1243,9 +1243,9 @@ impl DB for RocksDB {
 impl<'iter> DBIter<'iter> for RocksDB {
     type PrefixIter = PersistentPrefixIterator<'iter>;
 
-    fn iter_prefix(
+    fn iter_optional_prefix(
         &'iter self,
-        prefix: &Key,
+        prefix: Option<&Key>,
     ) -> PersistentPrefixIterator<'iter> {
         iter_subspace_prefix(self, prefix)
     }
@@ -1283,13 +1283,17 @@ impl<'iter> DBIter<'iter> for RocksDB {
 
 fn iter_subspace_prefix<'iter>(
     db: &'iter RocksDB,
-    prefix: &Key,
+    prefix: Option<&Key>,
 ) -> PersistentPrefixIterator<'iter> {
     let subspace_cf = db
         .get_column_family(SUBSPACE_CF)
         .expect("{SUBSPACE_CF} column family should exist");
     let db_prefix = "".to_owned();
-    iter_prefix(db, subspace_cf, db_prefix, prefix.to_string())
+    let prefix_string = match prefix {
+        Some(prefix) => prefix.to_string(),
+        None => "".to_string(),
+    };
+    iter_prefix(db, subspace_cf, db_prefix, prefix_string)
 }
 
 fn iter_diffs_prefix(

diff --git a/core/src/ledger/storage/mockdb.rs b/core/src/ledger/storage/mockdb.rs
@@ -514,9 +514,21 @@ impl DB for MockDB {
 impl<'iter> DBIter<'iter> for MockDB {
     type PrefixIter = MockPrefixIterator;
 
-    fn iter_prefix(&'iter self, prefix: &Key) -> MockPrefixIterator {
+    fn iter_optional_prefix(
+        &'iter self,
+        prefix: Option<&Key>,
+    ) -> MockPrefixIterator {
         let db_prefix = "subspace/".to_owned();
-        let prefix = format!("{}{}", db_prefix, prefix);
+        let prefix = format!(
+            "{}{}",
+            db_prefix,
+            match prefix {
+                None => "".to_string(),
+                Some(prefix) => {
+                    prefix.to_string()
+                }
+            }
+        );
         let iter = self.0.borrow().clone().into_iter();
         MockPrefixIterator::new(MockIterator { prefix, iter }, db_prefix)
     }

diff --git a/core/src/ledger/storage/mod.rs b/core/src/ledger/storage/mod.rs
@@ -340,7 +340,20 @@ pub trait DBIter<'iter> {
     ///
     /// Read account subspace key value pairs with the given prefix from the DB,
     /// ordered by the storage keys.
-    fn iter_prefix(&'iter self, prefix: &Key) -> Self::PrefixIter;
+    fn iter_prefix(&'iter self, prefix: &Key) -> Self::PrefixIter {
+        self.iter_optional_prefix(Some(prefix))
+    }
+
+    /// Iterate over all subspace keys
+    fn iter_all(&'iter self) -> Self::PrefixIter {
+        self.iter_optional_prefix(None)
+    }
+
+    /// Iterate over subspace keys, with optional prefix
+    fn iter_optional_prefix(
+        &'iter self,
+        prefix: Option<&Key>,
+    ) -> Self::PrefixIter;
 
     /// Read results subspace key value pairs from the DB
     fn iter_results(&'iter self) -> Self::PrefixIter;