feat: implement keys and values on StableBTreeMap

benchmarks/src/btreemap.rs

@@ -222,22 +222,62 @@ pub fn btreemap_insert_10mib_values() -> BenchResult {
 
 #[bench(raw)]
 pub fn btreemap_iter_small_values() -> BenchResult {
-    iter_helper(10_000, 0)
+    iter_helper(10_000, 0, IterType::Iter)
 }
 
 #[bench(raw)]
 pub fn btreemap_iter_rev_small_values() -> BenchResult {
-    iter_rev_helper(10_000, 0)
+    iter_helper(10_000, 0, IterType::IterRev)
 }
 
 #[bench(raw)]
 pub fn btreemap_iter_10mib_values() -> BenchResult {
-    iter_helper(200, 10 * 1024)
+    iter_helper(200, 10 * 1024, IterType::Iter)
 }
 
 #[bench(raw)]
 pub fn btreemap_iter_rev_10mib_values() -> BenchResult {
-    iter_rev_helper(200, 10 * 1024)
+    iter_helper(200, 10 * 1024, IterType::IterRev)
+}
+
+#[bench(raw)]
+pub fn btreemap_keys_small_values() -> BenchResult {
+    iter_helper(10_000, 0, IterType::Keys)
+}
+
+#[bench(raw)]
+pub fn btreemap_keys_rev_small_values() -> BenchResult {
+    iter_helper(10_000, 0, IterType::KeysRev)
+}
+
+#[bench(raw)]
+pub fn btreemap_keys_10mib_values() -> BenchResult {
+    iter_helper(200, 10 * 1024, IterType::Keys)
+}
+
+#[bench(raw)]
+pub fn btreemap_keys_rev_10mib_values() -> BenchResult {
+    iter_helper(200, 10 * 1024, IterType::KeysRev)
+}
+
+#[bench(raw)]
+pub fn btreemap_values_small_values() -> BenchResult {
+    iter_helper(10_000, 0, IterType::Values)
+}
+
+#[bench(raw)]
+pub fn btreemap_values_rev_small_values() -> BenchResult {
+    iter_helper(10_000, 0, IterType::ValuesRev)
+}
+
+#[bench(raw)]
+pub fn btreemap_values_10mib_values() -> BenchResult {
+    iter_helper(200, 10 * 1024, IterType::Values)
+}
+
+#[bench(raw)]
+pub fn btreemap_values_rev_10mib_values() -> BenchResult {
+    iter_helper(200, 10 * 1024, IterType::ValuesRev)
 }
 
 #[bench(raw)]
@@ -538,23 +578,20 @@ 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 {
+fn iter_helper(size: u32, value_size: u32, iter_type: IterType) -> 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]);
+    match iter_type {
+        IterType::Iter => bench_fn(|| for _ in btree.iter() {}),
+        IterType::IterRev => bench_fn(|| for _ in btree.iter().rev() {}),
+        IterType::Keys => bench_fn(|| for _ in btree.keys() {}),
+        IterType::KeysRev => bench_fn(|| for _ in btree.keys().rev() {}),
+        IterType::Values => bench_fn(|| for _ in btree.values() {}),
+        IterType::ValuesRev => bench_fn(|| for _ in btree.values().rev() {}),
     }
-
-    bench_fn(|| for _ in btree.iter().rev() {})
 }
 
 // Profiles getting a large number of random blobs from a btreemap.
@@ -630,3 +667,12 @@ fn remove_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
         }
     })
 }
+
+enum IterType {
+    Iter,
+    IterRev,
+    Keys,
+    KeysRev,
+    Values,
+    ValuesRev,
+}

canbench_results.yml

@@ -385,13 +385,13 @@ benches:
     scopes: {}
   btreemap_iter_rev_10mib_values:
     total:
-      instructions: 25585550
+      instructions: 25584039
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   btreemap_iter_rev_small_values:
     total:
-      instructions: 23878236
+      instructions: 23800315
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
@@ -401,6 +401,30 @@ benches:
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
+  btreemap_keys_10mib_values:
+    total:
+      instructions: 534290
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_keys_rev_10mib_values:
+    total:
+      instructions: 595467
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_keys_rev_small_values:
+    total:
+      instructions: 14369449
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_keys_small_values:
+    total:
+      instructions: 11184261
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
   btreemap_remove_blob_128_1024:
     total:
       instructions: 1916049094
@@ -533,6 +557,30 @@ benches:
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
+  btreemap_values_10mib_values:
+    total:
+      instructions: 17277830
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_values_rev_10mib_values:
+    total:
+      instructions: 17276742
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_values_rev_small_values:
+    total:
+      instructions: 22619653
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
+  btreemap_values_small_values:
+    total:
+      instructions: 22560352
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: {}
   memory_manager_baseline:
     total:
       instructions: 1176577052
@@ -541,13 +589,13 @@ benches:
     scopes: {}
   memory_manager_grow:
     total:
-      instructions: 350727867
+      instructions: 351687872
       heap_increase: 2
       stable_memory_increase: 32000
     scopes: {}
   memory_manager_overhead:
     total:
-      instructions: 1182141676
+      instructions: 1182143127
       heap_increase: 0
       stable_memory_increase: 8320
     scopes: {}

src/btreemap.rs

@@ -51,6 +51,7 @@
 mod allocator;
 mod iter;
 mod node;
+use crate::btreemap::iter::{IterInternal, KeysIter, ValuesIter};
 use crate::{
     storable::Bound as StorableBound,
     types::{Address, NULL},
@@ -1006,33 +1007,62 @@ where
 
     /// Returns an iterator over the entries of the map, sorted by key.
     pub fn iter(&self) -> Iter<K, V, M> {
-        Iter::new(self)
+        self.iter_internal().into()
     }
 
     /// Returns an iterator over the entries in the map where keys
     /// belong to the specified range.
     pub fn range(&self, key_range: impl RangeBounds<K>) -> Iter<K, V, M> {
+        self.range_internal(key_range).into()
+    }
+
+    /// 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 let Some((start_key, _)) = self.range(..bound).next_back() {
+            IterInternal::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded)).into()
+        } else {
+            IterInternal::null(self).into()
+        }
+    }
+
+    /// Returns an iterator over the keys of the map.
+    pub fn keys(&self) -> KeysIter<K, V, M> {
+        self.iter_internal().into()
+    }
+
+    /// Returns an iterator over the keys of the map which belong to the specified range.
+    pub fn keys_range(&self, key_range: impl RangeBounds<K>) -> KeysIter<K, V, M> {
+        self.range_internal(key_range).into()
+    }
+
+    /// Returns an iterator over the values of the map, sorted by key.
+    pub fn values(&self) -> ValuesIter<K, V, M> {
+        self.iter_internal().into()
+    }
+
+    /// Returns an iterator over the values of the map where keys
+    /// belong to the specified range.
+    pub fn values_range(&self, key_range: impl RangeBounds<K>) -> ValuesIter<K, V, M> {
+        self.range_internal(key_range).into()
+    }
+
+    fn iter_internal(&self) -> IterInternal<K, V, M> {
+        IterInternal::new(self)
+    }
+
+    fn range_internal(&self, key_range: impl RangeBounds<K>) -> IterInternal<K, V, M> {
         if self.root_addr == NULL {
             // Map is empty.
-            return Iter::null(self);
+            return IterInternal::null(self);
         }
 
         let range = (
             key_range.start_bound().cloned(),
             key_range.end_bound().cloned(),
         );
 
-        Iter::new_in_range(self, range)
-    }
-
-    /// 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 let Some((start_key, _)) = self.range(..bound).next_back() {
-            Iter::new_in_range(self, (Bound::Included(start_key), Bound::Unbounded))
-        } else {
-            Iter::null(self)
-        }
+        IterInternal::new_in_range(self, range)
     }
 
     // Merges one node (`source`) into another (`into`), along with a median entry.