revert iter for kv_store

rumos-io · Oct 10, 2024 · ca34f23 · ca34f23
1 parent fc1d888
commit ca34f23
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 20 deletions.
diff --git a/kv_store/src/range.rs b/kv_store/src/range.rs
@@ -1,14 +1,27 @@
-use std::borrow::Cow;
+use std::{borrow::Cow, ops::Bound};
 
 use database::Database;
-use trees::iavl::Range as TreeRange;
+use trees::iavl;
 
 use crate::utils::MergedRange;
 
 #[derive(Debug, Clone)]
 pub enum Range<'a, DB> {
     Merged(MergedRange<'a>),
-    Tree(TreeRange<'a, DB>),
+    Tree(iavl::Range<'a, DB>),
+    MergedRev(std::iter::Rev<MergedRange<'a>>),
+    TreeRev(iavl::RevRange<'a, DB, Vec<u8>, (Bound<Vec<u8>>, Bound<Vec<u8>>)>),
+}
+
+impl<'a, DB> Range<'a, DB> {
+    pub fn rev_iter(self) -> Range<'a, DB> {
+        match self {
+            Range::Merged(range) => Range::MergedRev(range.rev()),
+            Range::Tree(range) => Range::TreeRev(range.rev_iter()),
+            Range::MergedRev(rev) => Range::MergedRev(rev),
+            Range::TreeRev(rev_range) => Range::TreeRev(rev_range),
+        }
+    }
 }
 
 impl<'a, DB: Database> Iterator for Range<'a, DB> {
@@ -20,6 +33,10 @@ impl<'a, DB: Database> Iterator for Range<'a, DB> {
             Range::Tree(range) => range
                 .next()
                 .map(|(first, second)| (Cow::Owned(first), Cow::Owned(second))),
+            Range::MergedRev(range) => range.next(),
+            Range::TreeRev(range) => range
+                .next()
+                .map(|(first, second)| (Cow::Owned(first), Cow::Owned(second))),
         }
     }
 }
@@ -30,8 +47,22 @@ impl<'a, DB> From<MergedRange<'a>> for Range<'a, DB> {
     }
 }
 
-impl<'a, DB> From<TreeRange<'a, DB>> for Range<'a, DB> {
-    fn from(value: TreeRange<'a, DB>) -> Self {
+impl<'a, DB> From<iavl::Range<'a, DB>> for Range<'a, DB> {
+    fn from(value: iavl::Range<'a, DB>) -> Self {
         Self::Tree(value)
     }
 }
+
+impl<'a, DB> From<std::iter::Rev<MergedRange<'a>>> for Range<'a, DB> {
+    fn from(value: std::iter::Rev<MergedRange<'a>>) -> Self {
+        Self::MergedRev(value)
+    }
+}
+
+impl<'a, DB> From<iavl::RevRange<'a, DB, Vec<u8>, (Bound<Vec<u8>>, Bound<Vec<u8>>)>>
+    for Range<'a, DB>
+{
+    fn from(value: iavl::RevRange<'a, DB, Vec<u8>, (Bound<Vec<u8>>, Bound<Vec<u8>>)>) -> Self {
+        Self::TreeRev(value)
+    }
+}
diff --git a/trees/src/iavl/range/mod.rs b/trees/src/iavl/range/mod.rs
@@ -17,6 +17,23 @@ pub struct Range<'a, DB> {
     node_db: &'a NodeDB<DB>,
 }
 
+impl<'a, DB> Range<'a, DB> {
+    pub fn rev_iter(self) -> RevRange<'a, DB, Vec<u8>, (Bound<Vec<u8>>, Bound<Vec<u8>>)> {
+        let Self {
+            range,
+            delayed_nodes,
+            node_db,
+        } = self;
+
+        RevRange {
+            range: range,
+            delayed_nodes: delayed_nodes.into(),
+            node_db,
+            _marker: PhantomData,
+        }
+    }
+}
+
 impl<'a, DB: Database> Range<'a, DB> {
     pub(crate) fn new<R: RangeBounds<Vec<u8>>>(
         range: R,
@@ -33,21 +50,6 @@ impl<'a, DB: Database> Range<'a, DB> {
         }
     }
 
-    pub fn rev_iter(self) -> RevRange<'a, DB, Vec<u8>, (Bound<Vec<u8>>, Bound<Vec<u8>>)> {
-        let Self {
-            range,
-            delayed_nodes,
-            node_db,
-        } = self;
-
-        RevRange {
-            range: range,
-            delayed_nodes: delayed_nodes.into(),
-            node_db,
-            _marker: PhantomData,
-        }
-    }
-
     fn traverse(&mut self) -> Option<(Vec<u8>, Vec<u8>)> {
         let node = self.delayed_nodes.pop()?;