feat(iterators): allow key-only iteration (#2092)

## Linked Issues/PRs
<!-- List of related issues/PRs -->

- #1997
- #2076

## Description
<!-- List of detailed changes -->
Implements `iter_store_keys` and `iter_keys` appropriately for in memory
stores and rocksdb. The previous PR made them call it, and this one
changes the underlying implementation to avoid value allocations.

## Checklist
- [x] Breaking changes are clearly marked as such in the PR description
and changelog (no breaking changes)
- [x] New behavior is reflected in tests (tests which test this
behaviour under the hood still pass)
- [x] [The specification](https://github.com/FuelLabs/fuel-specs/)
matches the implemented behavior (no changes to spec)

### Before requesting review
- [x] I have reviewed the code myself
- [x] I have created follow-up issues caused by this PR and linked them
here

### After merging, notify other teams

[Add or remove entries as needed]

- [ ] [Rust SDK](https://github.com/FuelLabs/fuels-rs/)
- [ ] [Sway compiler](https://github.com/FuelLabs/sway/)
- [ ] [Platform
documentation](https://github.com/FuelLabs/devrel-requests/issues/new?assignees=&labels=new+request&projects=&template=NEW-REQUEST.yml&title=%5BRequest%5D%3A+)
(for out-of-organization contributors, the person merging the PR will do
this)
- [ ] Someone else?

---------

Co-authored-by: Hannes Karppila <[email protected]>
Co-authored-by: green <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -8,12 +8,14 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 - [2051](https://github.com/FuelLabs/fuel-core/pull/2051): Add support for AWS KMS signing for the PoA consensus module. The new key can be specified with `--consensus-aws-kms AWS_KEY_ARN`.
+- [2092](https://github.com/FuelLabs/fuel-core/pull/2092): Allow iterating by keys in rocksdb, and other storages.
 
 ### Changed
 
 #### Breaking
 - [2051](https://github.com/FuelLabs/fuel-core/pull/2051): Misdocumented `CONSENSUS_KEY` environ variable has been removed, use `CONSENSUS_KEY_SECRET` instead. Also raises MSRV to `1.79.0`.
 
+
 ## [Version 0.33.0]
 
 ### Added

crates/fuel-core/src/service/genesis/importer/import_task.rs

@@ -172,6 +172,7 @@ mod tests {
         },
         kv_store::{
             KVItem,
+            KeyItem,
             KeyValueInspect,
             Value,
         },
@@ -553,6 +554,16 @@ mod tests {
         ) -> BoxedIter<KVItem> {
             unimplemented!()
         }
+
+        fn iter_store_keys(
+            &self,
+            _: Self::Column,
+            _: Option<&[u8]>,
+            _: Option<&[u8]>,
+            _: IterDirection,
+        ) -> BoxedIter<KeyItem> {
+            unimplemented!()
+        }
     }
 
     impl TransactableStorage<BlockHeight> for BrokenTransactions {

crates/fuel-core/src/state.rs

@@ -34,6 +34,8 @@ pub mod iterable_key_value_view;
 pub mod key_value_view;
 #[cfg(feature = "rocksdb")]
 pub mod rocks_db;
+#[cfg(feature = "rocksdb")]
+pub mod rocks_db_key_iterator;
 
 pub type ColumnType<Description> = <Description as DatabaseDescription>::Column;
 
@@ -157,4 +159,28 @@ where
             core::iter::empty().into_boxed()
         }
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<fuel_core_storage::kv_store::KeyItem> {
+        // We cannot define iter_store_keys appropriately for the `ChangesIterator`,
+        // because we have to filter out the keys that were removed, which are
+        // marked as `WriteOperation::Remove` in the value
+        // copied as-is from the above function, but only to return keys
+        if let Some(tree) = self.changes.get(&column.id()) {
+            fuel_core_storage::iter::iterator(tree, prefix, start, direction)
+                .filter_map(|(key, value)| match value {
+                    WriteOperation::Insert(_) => Some(key.clone().into()),
+                    WriteOperation::Remove => None,
+                })
+                .map(Ok)
+                .into_boxed()
+        } else {
+            core::iter::empty().into_boxed()
+        }
+    }
 }

crates/fuel-core/src/state/data_source.rs

@@ -86,4 +86,14 @@ where
     ) -> BoxedIter<KVItem> {
         self.data.iter_store(column, prefix, start, direction)
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<fuel_core_storage::kv_store::KeyItem> {
+        self.data.iter_store_keys(column, prefix, start, direction)
+    }
 }

crates/fuel-core/src/state/generic_database.rs

@@ -138,6 +138,17 @@ where
     ) -> BoxedIter<KVItem> {
         self.storage.iter_store(column, prefix, start, direction)
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<fuel_core_storage::kv_store::KeyItem> {
+        self.storage
+            .iter_store_keys(column, prefix, start, direction)
+    }
 }
 
 impl<Storage> AsRef<Storage> for GenericDatabase<Storage> {

crates/fuel-core/src/state/historical_rocksdb.rs

@@ -416,6 +416,17 @@ where
         self.db
             .iter_store(Column::OriginalColumn(column), prefix, start, direction)
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<fuel_core_storage::kv_store::KeyItem> {
+        self.db
+            .iter_store_keys(Column::OriginalColumn(column), prefix, start, direction)
+    }
 }
 
 impl<Description> TransactableStorage<Description::Height>

crates/fuel-core/src/state/historical_rocksdb/view_at_height.rs

@@ -9,7 +9,10 @@ use crate::{
             deserialize,
             height_key,
         },
-        rocks_db::RocksDb,
+        rocks_db::{
+            KeyAndValue,
+            RocksDb,
+        },
     },
 };
 use fuel_core_storage::{
@@ -50,7 +53,7 @@ where
         let height_key = height_key(key, &self.height);
         let options = ReadOptions::default();
         let nearest_modification = read_history
-            ._iter_all(
+            .iterator::<KeyAndValue>(
                 Column::HistoricalDuplicateColumn(column),
                 options,
                 IteratorMode::From(&height_key, rocksdb::Direction::Forward),

crates/fuel-core/src/state/in_memory/memory_store.rs

@@ -15,12 +15,14 @@ use crate::{
 use fuel_core_storage::{
     iter::{
         iterator,
+        keys_iterator,
         BoxedIter,
         IntoBoxedIter,
         IterableStore,
     },
     kv_store::{
         KVItem,
+        KeyItem,
         KeyValueInspect,
         StorageColumn,
         Value,
@@ -99,6 +101,22 @@ where
 
         collection.into_iter().map(Ok)
     }
+
+    pub fn iter_all_keys(
+        &self,
+        column: Description::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> impl Iterator<Item = KeyItem> {
+        let lock = self.inner[column.as_usize()].lock().expect("poisoned");
+
+        let collection: Vec<_> = keys_iterator(&lock, prefix, start, direction)
+            .map(|key| key.to_vec())
+            .collect();
+
+        collection.into_iter().map(Ok)
+    }
 }
 
 impl<Description> KeyValueInspect for MemoryStore<Description>
@@ -129,6 +147,17 @@ where
     ) -> BoxedIter<KVItem> {
         self.iter_all(column, prefix, start, direction).into_boxed()
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<fuel_core_storage::kv_store::KeyItem> {
+        self.iter_all_keys(column, prefix, start, direction)
+            .into_boxed()
+    }
 }
 
 impl<Description> TransactableStorage<Description::Height> for MemoryStore<Description>

crates/fuel-core/src/state/in_memory/memory_view.rs

@@ -5,13 +5,15 @@ use crate::database::database_description::{
 use fuel_core_storage::{
     iter::{
         iterator,
+        keys_iterator,
         BoxedIter,
         IntoBoxedIter,
         IterDirection,
         IterableStore,
     },
     kv_store::{
         KVItem,
+        KeyItem,
         KeyValueInspect,
         StorageColumn,
         Value,
@@ -47,6 +49,20 @@ where
             .map(|(key, value)| (key.clone().into(), value.clone()))
             .map(Ok)
     }
+
+    pub fn iter_all_keys(
+        &self,
+        column: Description::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> impl Iterator<Item = KeyItem> + '_ {
+        let btree = &self.inner[column.as_usize()];
+
+        keys_iterator(btree, prefix, start, direction)
+            .map(|key| key.clone().into())
+            .map(Ok)
+    }
 }
 
 impl<Description> KeyValueInspect for MemoryView<Description>
@@ -73,4 +89,15 @@ where
     ) -> BoxedIter<KVItem> {
         self.iter_all(column, prefix, start, direction).into_boxed()
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<KeyItem> {
+        self.iter_all_keys(column, prefix, start, direction)
+            .into_boxed()
+    }
 }

crates/fuel-core/src/state/iterable_key_value_view.rs

@@ -6,6 +6,7 @@ use fuel_core_storage::{
     },
     kv_store::{
         KVItem,
+        KeyItem,
         KeyValueInspect,
         StorageColumn,
         Value,
@@ -85,4 +86,14 @@ where
     ) -> BoxedIter<KVItem> {
         self.0.iter_store(column, prefix, start, direction)
     }
+
+    fn iter_store_keys(
+        &self,
+        column: Self::Column,
+        prefix: Option<&[u8]>,
+        start: Option<&[u8]>,
+        direction: IterDirection,
+    ) -> BoxedIter<KeyItem> {
+        self.0.iter_store_keys(column, prefix, start, direction)
+    }
 }