EPROD-1117 pop_first and pop_last for btreemap

ic-stable-structures/src/structure/cache/btreemap.rs

@@ -63,6 +63,20 @@ where
         self.inner.remove(key)
     }
 
+    fn pop_first(&mut self) -> Option<(K, V)> {
+        let (k, v) = self.inner.pop_first()?;
+        self.cache.remove(&k);
+
+        Some((k, v))
+    }
+
+    fn pop_last(&mut self) -> Option<(K, V)> {
+        let (k, v) = self.inner.pop_last()?;
+        self.cache.remove(&k);
+
+        Some((k, v))
+    }
+
     fn len(&self) -> u64 {
         self.inner.len()
     }
@@ -222,6 +236,35 @@ mod tests {
         assert_eq!(None, map.get(&2));
     }
 
+    #[test]
+    fn test_should_pop_first() {
+        let mut map = CachedStableBTreeMap::new(VectorMemory::default(), 10);
+
+        map.insert(0u32, 42u32);
+        map.insert(10, 100);
+
+        assert_eq!(map.pop_first(), Some((0, 42)));
+
+        // try to get
+        assert!(map.get(&0).is_none());
+
+        assert_eq!(map.len(), 1);
+    }
+
+    #[test]
+    fn test_should_pop_last() {
+        let mut map = CachedStableBTreeMap::new(VectorMemory::default(), 10);
+
+        map.insert(0u32, 42u32);
+        map.insert(10, 100);
+
+        assert_eq!(map.pop_last(), Some((10, 100)));
+
+        assert!(map.get(&10).is_none());
+
+        assert_eq!(map.len(), 1);
+    }
+
     #[test]
     fn should_replace_old_value() {
         let cache_items = 2;

ic-stable-structures/src/structure/mod.rs

@@ -27,6 +27,12 @@ pub trait BTreeMapStructure<K, V> {
     ///   - `key.to_bytes().len() <= K::MAX_SIZE`
     fn remove(&mut self, key: &K) -> Option<V>;
 
+    /// Removes and returns the first element in the map.
+    fn pop_first(&mut self) -> Option<(K, V)>;
+
+    /// Removes and returns the last element in the map.
+    fn pop_last(&mut self) -> Option<(K, V)>;
+
     /// True if contains the key.
     fn contains_key(&self, key: &K) -> bool;
 

ic-stable-structures/src/structure/stable_storage/btreemap.rs

@@ -46,6 +46,14 @@ where
         self.0.remove(key)
     }
 
+    fn pop_first(&mut self) -> Option<(K, V)> {
+        self.0.pop_first()
+    }
+
+    fn pop_last(&mut self) -> Option<(K, V)> {
+        self.0.pop_last()
+    }
+
     fn len(&self) -> u64 {
         self.0.len()
     }
@@ -204,6 +212,32 @@ mod tests {
         assert_eq!(map.len(), 2);
     }
 
+    #[test]
+    fn test_should_pop_first() {
+        let mut map = StableBTreeMap::new(VectorMemory::default());
+
+        map.insert(0u32, 42u32);
+        map.insert(10, 100);
+
+        assert_eq!(map.pop_first(), Some((0, 42)));
+
+        assert_eq!(map.len(), 1);
+        assert!(map.get(&0).is_none());
+    }
+
+    #[test]
+    fn test_should_pop_last() {
+        let mut map = StableBTreeMap::new(VectorMemory::default());
+
+        map.insert(0u32, 42u32);
+        map.insert(10, 100);
+
+        assert_eq!(map.pop_last(), Some((10, 100)));
+
+        assert_eq!(map.len(), 1);
+        assert!(map.get(&10).is_none());
+    }
+
     #[test]
     fn iter_test() {
         let mut map = StableBTreeMap::new(VectorMemory::default());

ic-task-scheduler/tests/pocket_ic_tests/mod.rs

@@ -69,7 +69,7 @@ impl PocketIcTestContext {
     }
 }
 
-async fn deploy_dummy_scheduler_canister<'a>() -> anyhow::Result<PocketIcTestContext> {
+async fn deploy_dummy_scheduler_canister() -> anyhow::Result<PocketIcTestContext> {
     let client = init_pocket_ic().await.build_async().await;
     println!("Creating dummy canister");