Refactor iter_store_keys

crates/storage/src/iter/changes_iterator.rs

@@ -14,8 +14,13 @@ use crate::{
         Value,
         WriteOperation,
     },
-    transactional::StorageChanges,
+    transactional::{
+        Changes,
+        ReferenceBytesKey,
+        StorageChanges,
+    },
 };
+use alloc::collections::BTreeMap;
 
 /// A type that allows to iterate over the `StorageChanges`.
 pub struct ChangesIterator<'a, Column> {
@@ -137,19 +142,13 @@ where
         // marked as `WriteOperation::Remove` in the value
         // copied as-is from the above function, but only to return keys
         match self.changes {
-            StorageChanges::Changes(changes) => {
-                if let Some(tree) = changes.get(&column.id()) {
-                    crate::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()
-                }
-            }
+            StorageChanges::Changes(changes) => get_insert_keys_from_changes(
+                changes,
+                column.id(),
+                prefix,
+                start,
+                direction,
+            ),
             StorageChanges::ChangesList(changes_list) => {
                 // We have to clone the prefix and start, because we need to pass them to the iterator
                 // if someone finds a solution without making it a vec, feel free to contribute :)
@@ -158,24 +157,46 @@ where
                 let start = start.map(|start| start.to_vec());
                 changes_list
                     .iter()
-                    .filter_map(move |changes| {
-                        changes.get(&column).map(|tree| {
-                            crate::iter::iterator(
-                                tree,
-                                prefix.as_deref(),
-                                start.as_deref(),
-                                direction,
-                            )
-                            .filter_map(|(key, value)| match value {
-                                WriteOperation::Insert(_) => Some(key.clone().into()),
-                                WriteOperation::Remove => None,
-                            })
-                            .map(Ok)
-                        })
+                    .flat_map(move |changes| {
+                        get_insert_keys_from_changes(
+                            changes,
+                            column,
+                            prefix.as_deref(),
+                            start.as_deref(),
+                            direction,
+                        )
                     })
-                    .flatten()
                     .into_boxed()
             }
         }
     }
 }
+
+fn get_insert_keys_from_changes<'a>(
+    changes: &'a Changes,
+    column_id: u32,
+    prefix: Option<&[u8]>,
+    start: Option<&[u8]>,
+    direction: IterDirection,
+) -> BoxedIter<'a, crate::kv_store::KeyItem> {
+    changes
+        .get(&column_id)
+        .map_or(core::iter::empty().into_boxed(), |tree| {
+            boxed_insert_iter(tree, prefix, start, direction)
+        })
+}
+
+fn boxed_insert_iter<'a>(
+    tree: &'a BTreeMap<ReferenceBytesKey, WriteOperation>,
+    prefix: Option<&[u8]>,
+    start: Option<&[u8]>,
+    direction: IterDirection,
+) -> BoxedIter<'a, crate::kv_store::KeyItem> {
+    crate::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()
+}