feat: implement DoubleEndedIterator for StableBTreeMap

benchmarks/src/btreemap.rs

@@ -220,6 +220,26 @@ pub fn btreemap_insert_10mib_values() -> BenchResult {
     })
 }
 
+#[bench(raw)]
+pub fn btreemap_iter_small_values() -> BenchResult {
+    iter_helper(10_000, 0)
+}
+
+#[bench(raw)]
+pub fn btreemap_iter_rev_small_values() -> BenchResult {
+    iter_rev_helper(10_000, 0)
+}
+
+#[bench(raw)]
+pub fn btreemap_iter_10mib_values() -> BenchResult {
+    iter_helper(200, 10 * 1024)
+}
+
+#[bench(raw)]
+pub fn btreemap_iter_rev_10mib_values() -> BenchResult {
+    iter_rev_helper(200, 10 * 1024)
+}
+
 #[bench(raw)]
 pub fn btreemap_iter_count_small_values() -> BenchResult {
     let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
@@ -517,6 +537,26 @@ fn insert_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
     })
 }
 
+// Profiles iterating over a btreemap.
+fn iter_helper(size: u32, value_size: u32) -> BenchResult {
+    let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
+    for i in 0..size {
+        btree.insert(i, vec![0u8; value_size as usize]);
+    }
+
+    bench_fn(|| for _ in btree.iter() {})
+}
+
+// Profiles iterating in reverse over a btreemap.
+fn iter_rev_helper(size: u32, value_size: u32) -> BenchResult {
+    let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
+    for i in 0..size {
+        btree.insert(i, vec![0u8; value_size as usize]);
+    }
+
+    bench_fn(|| for _ in btree.iter().rev() {})
+}
+
 // Profiles getting a large number of random blobs from a btreemap.
 fn get_blob_helper<const K: usize, const V: usize>() -> BenchResult {
     let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());

canbench_results.yml

@@ -365,15 +365,39 @@ benches:
       heap_increase: 0
       stable_memory_increase: 6
     scopes: {}
+  btreemap_iter_10mib_values:
+    total:
+      instructions: 25583733
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
   btreemap_iter_count_10mib_values:
     total:
-      instructions: 525036
+      instructions: 544088
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_iter_count_small_values:
     total:
-      instructions: 10458516
+      instructions: 11007833
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_iter_rev_10mib_values:
+    total:
+      instructions: 25585550
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_iter_rev_small_values:
+    total:
+      instructions: 23878236
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_iter_small_values:
+    total:
+      instructions: 23721014
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -517,13 +541,13 @@ benches:
     scopes: {}
   memory_manager_grow:
     total:
-      instructions: 351687872
+      instructions: 350727867
       heap_increase: 2
       stable_memory_increase: 32000
     scopes: {}
   memory_manager_overhead:
     total:
-      instructions: 1182143127
+      instructions: 1182141676
       heap_increase: 0
       stable_memory_increase: 8320
     scopes: {}

src/btreemap.rs

@@ -58,7 +58,6 @@ use crate::{
 };
 use allocator::Allocator;
 pub use iter::Iter;
-use iter::{Cursor, Index};
 use node::{DerivedPageSize, Entry, Node, NodeType, PageSize, Version};
 use std::borrow::Cow;
 use std::marker::PhantomData;
@@ -1029,74 +1028,10 @@ where
     /// Returns an iterator pointing to the first element below the given bound.
     /// Returns an empty iterator if there are no keys below the given bound.
     pub fn iter_upper_bound(&self, bound: &K) -> Iter<K, V, M> {
-        if self.root_addr == NULL {
-            // Map is empty.
-            return Iter::null(self);
-        }
-
-        let dummy_bounds = (Bound::Unbounded, Bound::Unbounded);
-        // INVARIANT: all cursors point to keys greater than or equal to bound.
-        let mut cursors = vec![];
-
-        let mut node = self.load_node(self.root_addr);
-        loop {
-            match node.search(bound) {
-                Ok(idx) | Err(idx) => {
-                    match node.node_type() {
-                        NodeType::Leaf => {
-                            if idx == 0 {
-                                // We descended into a leaf but didn't find a node less than
-                                // the upper bound. Thus we unwind the cursor stack until we
-                                // hit a cursor pointing to an element other than the first key,
-                                // and we shift the position backward. If there is no such cursor,
-                                // the bound must be <= min element, so we return an empty iterator.
-                                while let Some(cursor) = cursors.pop() {
-                                    match cursor {
-                                        Cursor::Node {
-                                            node,
-                                            next: Index::Entry(n),
-                                        } => {
-                                            if n == 0 {
-                                                debug_assert!(node.key(n) >= bound);
-                                                continue;
-                                            } else {
-                                                debug_assert!(node.key(n - 1) < bound);
-                                                cursors.push(Cursor::Node {
-                                                    node,
-                                                    next: Index::Entry(n - 1),
-                                                });
-                                                break;
-                                            }
-                                        }
-                                        _ => panic!("BUG: unexpected cursor shape"),
-                                    }
-                                }
-                                // If the cursors are empty, the iterator will be empty.
-                                return Iter::new_with_cursors(self, dummy_bounds, cursors);
-                            }
-                            debug_assert!(node.key(idx - 1) < bound);
-
-                            cursors.push(Cursor::Node {
-                                node,
-                                next: Index::Entry(idx - 1),
-                            });
-                            return Iter::new_with_cursors(self, dummy_bounds, cursors);
-                        }
-                        NodeType::Internal => {
-                            let child = self.load_node(node.child(idx));
-                            // We push the node even if idx == node.entries_len()
-                            // If we find the position in the child, the iterator will skip this
-                            // cursor. But if the all keys in the child are greater than or equal to
-                            // the bound, we will be able to use this cursor as a fallback.
-                            cursors.push(Cursor::Node {
-                                node,
-                                next: Index::Entry(idx),
-                            });
-                            node = child;
-                        }
-                    }
-                }
-            }
+        if let Some((start_key, _)) = self.range(..bound).next_back() {
+            Iter::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded))
+        } else {
+            Iter::null(self)
         }
     }