Merge pull request #20 from seniorjoinu/develop

Develop

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ic-stable-memory"
-version = "0.4.1"
+version = "0.4.4"
 authors = ["Александр Втюрин <[email protected]>"]
 edition = "2021"
 description = "Internet Computer's stable memory collections and tools"
@@ -12,10 +12,10 @@ keywords = ["dfinity", "internet-computer", "ic", "stable-memory", "collections"
 path = "./src/lib.rs"
 
 [dependencies]
-ic-cdk = "0.6.10"
+ic-cdk = "0.7.3"
 candid = "0.8.4"
 serde = "1.0.152"
-serde_bytes = "0.11.8"
+serde_bytes = "0.11.9"
 num-bigint = "0.4.3"
 sha2 = "0.10.6"
 zwohash = "0.1.2"
@@ -24,7 +24,7 @@ ic-stable-memory-derive = "0.4.2"
 [dev-dependencies]
 rand = "0.8.5"
 ic-cdk-macros = "0.6.8"
-candid_derive = "0.5.0"
+candid_derive = "0.6.0"
 ic-certified-map = "0.3.2"
 serde_test = "1.0.152"
 

src/collections/btree_map/iter.rs

@@ -32,27 +32,28 @@ impl<'a, K: StableType + AsFixedSizeBytes + Ord, V: StableType + AsFixedSizeByte
 
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(node) = &self.node {
-            let k = node.get_key(self.node_idx);
-            let v = node.get_value(self.node_idx);
-
-            if self.node_idx == self.node_len - 1 {
-                let next_ptr = u64::from_fixed_size_bytes(&node.read_next_ptr_buf());
+            if self.node_idx == self.node_len {
+                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);
-            } else {
-                self.node_idx += 1;
             }
 
-            Some((k, v))
+            let res = (&self.node)
+                .as_ref()
+                .map(|it| (it.get_key(self.node_idx), it.get_value(self.node_idx)));
+
+            self.node_idx += 1;
+
+            res
         } else {
             let mut node = unsafe { self.root.as_ref()?.copy() };
             let leaf = loop {
@@ -80,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) {
@@ -1720,8 +1742,10 @@ impl LeveledList {
         match self {
             LeveledList::None => {}
             LeveledList::Some((v, _)) => {
-                if let Ok(idx) = v[level].binary_search(&ptr) {
-                    v[level].remove(idx);
+                if let Some(level_list) = v.get_mut(level) {
+                    if let Ok(idx) = level_list.binary_search(&ptr) {
+                        level_list.remove(idx);
+                    }
                 }
             }
         }
@@ -1731,7 +1755,8 @@ impl LeveledList {
         match self {
             LeveledList::None => unreachable!(),
             LeveledList::Some((v, max_level)) => {
-                let mut ptr = v[*max_level].pop();
+                let level_list = v.get_mut(*max_level)?;
+                let mut ptr = level_list.pop();
 
                 while ptr.is_none() {
                     if *max_level == 0 {
@@ -1925,10 +1950,14 @@ mod tests {
                 assert_eq!(i, *k);
                 assert_eq!(i, *v);
 
+                print!("({:?}, {:?}), ", *k, *v);
+
                 i += 1;
             }
 
-            assert_eq!(i, 199);
+            println!();
+
+            assert_eq!(i, 200);
         }
 
         _debug_validate_allocator();
@@ -2061,6 +2090,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> {