add certified btree set, add get_random

src/collections/btree_map/iter.rs

@@ -33,28 +33,27 @@ impl<'a, K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeByte
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(node) = &self.node {
             if self.node_idx == self.node_len {
-                let next_ptr = u64::from_fixed_size_bytes(&node.read_next_ptr_buf());
+                let ptr = u64::from_fixed_size_bytes(&node.read_next_ptr_buf());
 
-                if next_ptr == 0 {
+                if ptr == 0 {
                     return None;
                 }
 
-                let next = unsafe { LeafBTreeNode::<K, V>::from_ptr(next_ptr) };
-                let len = next.read_len();
+                let new_node = unsafe { LeafBTreeNode::<K, V>::from_ptr(ptr) };
+                let len = new_node.read_len();
 
+                self.node = Some(new_node);
                 self.node_idx = 0;
                 self.node_len = len;
-                self.node = Some(next);
+            }
 
-                self.next()
-            } else {
-                let k = node.get_key(self.node_idx);
-                let v = node.get_value(self.node_idx);
+            let res = (&self.node)
+                .as_ref()
+                .map(|it| (it.get_key(self.node_idx), it.get_value(self.node_idx)));
 
-                self.node_idx += 1;
+            self.node_idx += 1;
 
-                Some((k, v))
-            }
+            res
         } else {
             let mut node = unsafe { self.root.as_ref()?.copy() };
             let leaf = loop {
@@ -82,3 +81,60 @@ impl<'a, K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeByte
         }
     }
 }
+
+impl<'a, K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeBytes>
+    DoubleEndedIterator for SBTreeMapIter<'a, K, V>
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if let Some(node) = &self.node {
+            if self.node_idx == 0 {
+                return None;
+            }
+
+            self.node_idx -= 1;
+
+            let k = node.get_key(self.node_idx);
+            let v = node.get_value(self.node_idx);
+
+            if self.node_idx == 0 {
+                let ptr = u64::from_fixed_size_bytes(&node.read_prev_ptr_buf());
+
+                if ptr != 0 {
+                    let new_node = unsafe { LeafBTreeNode::<K, V>::from_ptr(ptr) };
+                    let len = new_node.read_len();
+
+                    self.node = Some(new_node);
+                    self.node_idx = len;
+                    self.node_len = len;
+                }
+            }
+
+            Some((k, v))
+        } else {
+            let mut node = unsafe { self.root.as_ref()?.copy() };
+            let leaf = loop {
+                match node {
+                    BTreeNode::Internal(i) => {
+                        let len = i.read_len();
+                        let child_ptr = u64::from_fixed_size_bytes(&i.read_child_ptr_buf(len));
+                        node = BTreeNode::<K, V>::from_ptr(child_ptr);
+                    }
+                    BTreeNode::Leaf(l) => {
+                        break l;
+                    }
+                }
+            };
+
+            self.node_len = leaf.read_len();
+
+            if self.node_len == 0 {
+                return None;
+            }
+
+            self.node_idx = self.node_len;
+            self.node = Some(leaf);
+
+            self.next_back()
+        }
+    }
+}

src/collections/btree_map/mod.rs

@@ -8,6 +8,7 @@ use crate::mem::{StablePtr, StablePtrBuf};
 use crate::primitive::s_ref::SRef;
 use crate::primitive::s_ref_mut::SRefMut;
 use crate::primitive::StableType;
+use crate::utils::math::shuffle_bits;
 use crate::{isoprint, make_sure_can_allocate, OutOfMemory, SSlice};
 use std::borrow::Borrow;
 use std::fmt::{Debug, Formatter};
@@ -346,6 +347,42 @@ impl<K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeBytes> S
         Some(leaf_node.get_value(idx))
     }
 
+    /// Returns a random key, deterministically deriving the randomness from the seed.
+    /// This function is usefull, when you have a source of real randomness.
+    ///
+    /// If the collection is empty, returns [None].
+    ///
+    /// Same seed on the same collection leads to the same returned key.
+    /// Same seed on a modified collection may still lead to the same returned key.
+    /// You can use [utils::math::shuffle_bits] function to pseudo-randomly generate more seeds.
+    pub fn get_random_key(&self, mut seed: u32) -> Option<SRef<K>> {
+        if self.is_empty() {
+            return None;
+        }
+
+        let mut node = self.get_root()?;
+
+        loop {
+            match node {
+                BTreeNode::Internal(i) => {
+                    let len = i.read_len();
+                    let idx = seed as usize % (len + 1);
+                    let child_ptr = u64::from_fixed_size_bytes(&i.read_child_ptr_buf(idx));
+
+                    seed = shuffle_bits(seed);
+
+                    node = BTreeNode::from_ptr(child_ptr);
+                }
+                BTreeNode::Leaf(l) => {
+                    let len = l.read_len();
+                    let idx = seed as usize % len;
+
+                    break Some(l.get_key(idx));
+                }
+            }
+        }
+    }
+
     /// Returns a mutable reference [SRefMut] to a value stored by the key
     ///
     /// See also [SBTreeMap::get].
@@ -450,6 +487,33 @@ impl<K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeBytes> S
     ///
     ///     i += 1;
     /// }
+    ///
+    /// assert_eq!(i, 100);
+    /// ```
+    ///
+    /// One can use `.rev()` to get elements in reverse order.
+    ///
+    /// # Example
+    /// ```rust
+    /// # use ic_stable_memory::collections::SBTreeMap;
+    /// # use ic_stable_memory::stable_memory_init;
+    /// # unsafe { ic_stable_memory::mem::clear(); }
+    /// # stable_memory_init();
+    /// let mut map = SBTreeMap::new();
+    ///
+    /// for i in 0..100 {
+    ///     map.insert(i, i).expect("Out of memory");
+    /// }
+    ///
+    /// let mut i = 100;
+    /// for (k, v) in map.iter().rev() {
+    ///     i -= 1;
+    ///
+    ///     assert_eq!(*k, i);
+    ///     assert_eq!(*v, i);
+    /// }
+    ///
+    /// assert_eq!(i, 0);
     /// ```
     #[inline]
     pub fn iter(&self) -> SBTreeMapIter<K, V> {
@@ -468,48 +532,6 @@ impl<K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeBytes> S
         self.len() == 0
     }
 
-    /// Returns an immutable reference [SRef] to the first key-value pair in order
-    ///
-    /// If the collection is empty, returns [None]
-    pub fn first(&self) -> Option<(SRef<K>, SRef<V>)> {
-        let mut node = self.get_root()?;
-
-        loop {
-            match node {
-                BTreeNode::Internal(n) => {
-                    let ptr = u64::from_fixed_size_bytes(&n.read_child_ptr_buf(0));
-                    node = BTreeNode::from_ptr(ptr);
-                }
-                BTreeNode::Leaf(n) => {
-                    return Some((n.get_key(0), n.get_value(0)));
-                }
-            }
-        }
-    }
-
-    /// Returns a mutable reference [SRefMut] to the first key-value pair in order
-    ///
-    /// If the collection is empty, returns [None]
-    pub fn last(&self) -> Option<(SRef<K>, SRef<V>)> {
-        let mut node = self.get_root()?;
-
-        loop {
-            match node {
-                BTreeNode::Internal(n) => {
-                    let len = n.read_len();
-
-                    let ptr = u64::from_fixed_size_bytes(&n.read_child_ptr_buf(len));
-                    node = BTreeNode::from_ptr(ptr);
-                }
-                BTreeNode::Leaf(n) => {
-                    let len = n.read_len();
-
-                    return Some((n.get_key(len - 1), n.get_value(len - 1)));
-                }
-            }
-        }
-    }
-
     /// Removes all key-value pairs from this collection, releasing all occupied stable memory
     #[inline]
     pub fn clear(&mut self) {
@@ -2065,6 +2087,16 @@ mod tests {
             _debug_validate_allocator();
             assert_eq!(self.map().len() as usize, self.example.len());
 
+            // check random key
+            let seed: u32 = self.rng.gen();
+            let rand_key = self.map.as_ref().unwrap().get_random_key(seed);
+            if self.map.as_ref().unwrap().is_empty() {
+                assert!(rand_key.is_none());
+            } else {
+                assert!(rand_key.is_some());
+            }
+
+            // check consistency
             for key in self.keys.clone() {
                 let contains = self.map().contains_key(&key);
                 assert!(contains);

src/collections/btree_set/iter.rs

@@ -23,3 +23,9 @@ impl<'a, T: StableType + AsFixedSizeBytes + Ord> Iterator for SBTreeSetIter<'a,
         self.iter.next().map(|it| it.0)
     }
 }
+
+impl<'a, T: StableType + AsFixedSizeBytes + Ord> DoubleEndedIterator for SBTreeSetIter<'a, T> {
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.next_back().map(|it| it.0)
+    }
+}

src/collections/btree_set/mod.rs

@@ -1,6 +1,7 @@
 use crate::collections::btree_map::SBTreeMap;
 use crate::collections::btree_set::iter::SBTreeSetIter;
 use crate::encoding::AsFixedSizeBytes;
+use crate::primitive::s_ref::SRef;
 use crate::primitive::StableType;
 use std::borrow::Borrow;
 use std::fmt::{Debug, Formatter};
@@ -70,6 +71,12 @@ impl<T: Ord + StableType + AsFixedSizeBytes> SBTreeSet<T> {
         self.map.contains_key(value)
     }
 
+    /// See [SBTreeMap::get_random_key]
+    #[inline]
+    pub fn get_random(&self, seed: u32) -> Option<SRef<T>> {
+        self.map.get_random_key(seed)
+    }
+
     /// See [SBTreeMap::iter]
     #[inline]
     pub fn iter(&self) -> SBTreeSetIter<T> {

src/collections/certified_btree_map/mod.rs

@@ -187,7 +187,7 @@ pub struct SCertifiedBTreeMap<
     K: StableType + AsFixedSizeBytes + Ord + AsHashableBytes,
     V: StableType + AsFixedSizeBytes + AsHashTree,
 > {
-    inner: SBTreeMap<K, V>,
+    pub(crate) inner: SBTreeMap<K, V>,
     modified: LeveledList,
     uncommited: bool,
 }
@@ -293,6 +293,12 @@ impl<
         self.inner.get(key)
     }
 
+    /// See [SBTreeMap::get]
+    #[inline]
+    pub fn get_random_key(&self, seed: u32) -> Option<SRef<K>> {
+        self.inner.get_random_key(seed)
+    }
+
     /// Allows mutation of the value stored by the provided key, accepting a lambda to perform it
     ///
     /// This method recomputes the underlying Merkle tree, if the key-value pair is found
@@ -343,18 +349,6 @@ impl<
         self.inner.iter()
     }
 
-    /// See [SBTreeMap::first]
-    #[inline]
-    pub fn first(&self) -> Option<(SRef<K>, SRef<V>)> {
-        self.inner.first()
-    }
-
-    /// See [SBTreeMap::last]
-    #[inline]
-    pub fn last(&self) -> Option<(SRef<K>, SRef<V>)> {
-        self.inner.last()
-    }
-
     /// Commits all `uncommited` changes to this data structure, recalculating the underlying Merkle
     /// tree
     ///
@@ -520,8 +514,8 @@ impl<
     fn hash_tree(&self) -> HashTree {
         assert!(!self.uncommited);
 
-        let e1 = self.inner.first();
-        let e2 = self.inner.last();
+        let e1 = self.inner.iter().next();
+        let e2 = self.inner.iter().rev().next();
 
         match (e1, e2) {
             (None, None) => HashTree::Empty,

src/collections/certified_btree_set/iter.rs

@@ -0,0 +1,34 @@
+use crate::collections::btree_map::iter::SBTreeMapIter;
+use crate::collections::certified_btree_set::SCertifiedBTreeSet;
+use crate::encoding::AsFixedSizeBytes;
+use crate::primitive::s_ref::SRef;
+use crate::primitive::StableType;
+use crate::AsHashableBytes;
+
+pub struct SCertifiedBTreeSetIter<'a, T> {
+    iter: SBTreeMapIter<'a, T, ()>,
+}
+
+impl<'a, T: StableType + AsFixedSizeBytes + Ord + AsHashableBytes> SCertifiedBTreeSetIter<'a, T> {
+    pub fn new(set: &'a SCertifiedBTreeSet<T>) -> Self {
+        Self {
+            iter: SBTreeMapIter::new(&set.map.inner),
+        }
+    }
+}
+
+impl<'a, T: StableType + AsFixedSizeBytes + Ord> Iterator for SCertifiedBTreeSetIter<'a, T> {
+    type Item = SRef<'a, T>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter.next().map(|it| it.0)
+    }
+}
+
+impl<'a, T: StableType + AsFixedSizeBytes + Ord> DoubleEndedIterator
+    for SCertifiedBTreeSetIter<'a, T>
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.next_back().map(|it| it.0)
+    }
+}