From 7aaf3cea728ab4e104987b630ab359191eb179f1 Mon Sep 17 00:00:00 2001
From: "Raymond E. Pasco" <raymond@heliax.dev>
Date: Tue, 23 May 2023 11:20:01 -0400
Subject: [PATCH] storage: add optional prefix iterator methods

Subspace key iterators are created with the iter_prefix() method on the
storage trait, which takes a string as a key prefix. Because the empty
string should be an invalid key for parsing purposes, other topics have
changed key parsing to error on the empty string. However, this leaves
us without a way to iterate over all subspace keys, previously achieved
with iter_prefix("").

To remedy this, iter_prefix has been extended into
iter_optional_prefix, taking an optional Key. Calls to iter_prefix("")
can be replaced with iter_all(), a new method; iter_prefix() still
exists with its original signature but will no longer accept empty
strings.
---
 apps/src/lib/node/ledger/storage/rocksdb.rs | 12 ++++++++----
 core/src/ledger/storage/mockdb.rs           | 16 ++++++++++++++--
 core/src/ledger/storage/mod.rs              | 15 ++++++++++++++-
 3 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/apps/src/lib/node/ledger/storage/rocksdb.rs b/apps/src/lib/node/ledger/storage/rocksdb.rs
index 9d54bc6de3..dd7c534c3a 100644
--- a/apps/src/lib/node/ledger/storage/rocksdb.rs
+++ b/apps/src/lib/node/ledger/storage/rocksdb.rs
@@ -1188,9 +1188,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)
     }
@@ -1228,13 +1228,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
index 16e28d2759..a5db20582c 100644
--- a/core/src/ledger/storage/mockdb.rs
+++ b/core/src/ledger/storage/mockdb.rs
@@ -480,9 +480,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
index cb3f9299df..87e6edf6c1 100644
--- a/core/src/ledger/storage/mod.rs
+++ b/core/src/ledger/storage/mod.rs
@@ -321,7 +321,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;