Merge branch 'main' into main

CosmWasm · Sep 27, 2022 · 9b051c4 · 9b051c4
2 parents bd08192 + 6b7c704
commit 9b051c4
Show file tree

Hide file tree

Showing 2 changed files with 119 additions and 0 deletions.
diff --git a/packages/storage-plus/src/indexed_map.rs b/packages/storage-plus/src/indexed_map.rs
@@ -139,6 +139,34 @@ where
     fn no_prefix_raw(&self) -> Prefix<Vec<u8>, T, K> {
         Prefix::new(self.pk_namespace, &[])
     }
+
+    /// Clears the map, removing all elements.
+    pub fn clear(&self, store: &mut dyn Storage) {
+        const TAKE: usize = 10;
+        let mut cleared = false;
+
+        while !cleared {
+            let paths = self
+                .no_prefix_raw()
+                .keys_raw(store, None, None, cosmwasm_std::Order::Ascending)
+                .map(|raw_key| Path::<T>::new(self.pk_namespace, &[raw_key.as_slice()]))
+                // Take just TAKE elements to prevent possible heap overflow if the Map is big.
+                .take(TAKE)
+                .collect::<Vec<_>>();
+
+            paths.iter().for_each(|path| store.remove(path));
+
+            cleared = paths.len() < TAKE;
+        }
+    }
+
+    /// Returns `true` if the map is empty.
+    pub fn is_empty(&self, store: &dyn Storage) -> bool {
+        self.no_prefix_raw()
+            .keys_raw(store, None, None, cosmwasm_std::Order::Ascending)
+            .next()
+            .is_none()
+    }
 }
 
 #[cfg(feature = "iterator")]
@@ -1679,4 +1707,29 @@ mod test {
             );
         }
     }
+
+    #[test]
+    fn clear_works() {
+        let mut storage = MockStorage::new();
+        let map = build_map();
+        let (pks, _) = save_data(&mut storage, &map);
+
+        map.clear(&mut storage);
+
+        for key in pks {
+            assert!(!map.has(&storage, key));
+        }
+    }
+
+    #[test]
+    fn is_empty_works() {
+        let mut storage = MockStorage::new();
+        let map = build_map();
+
+        assert!(map.is_empty(&storage));
+
+        save_data(&mut storage, &map);
+
+        assert!(!map.is_empty(&storage));
+    }
 }
diff --git a/packages/storage-plus/src/map.rs b/packages/storage-plus/src/map.rs
@@ -109,6 +109,36 @@ where
             from_slice(&result).map(Some)
         }
     }
+
+    /// Clears the map, removing all elements.
+    #[cfg(feature = "iterator")]
+    pub fn clear(&self, store: &mut dyn Storage) {
+        const TAKE: usize = 10;
+        let mut cleared = false;
+
+        while !cleared {
+            let paths = self
+                .no_prefix_raw()
+                .keys_raw(store, None, None, cosmwasm_std::Order::Ascending)
+                .map(|raw_key| Path::<T>::new(self.namespace, &[raw_key.as_slice()]))
+                // Take just TAKE elements to prevent possible heap overflow if the Map is big.
+                .take(TAKE)
+                .collect::<Vec<_>>();
+
+            paths.iter().for_each(|path| store.remove(path));
+
+            cleared = paths.len() < TAKE;
+        }
+    }
+
+    /// Returns `true` if the map is empty.
+    #[cfg(feature = "iterator")]
+    pub fn is_empty(&self, store: &dyn Storage) -> bool {
+        self.no_prefix_raw()
+            .keys_raw(store, None, None, cosmwasm_std::Order::Ascending)
+            .next()
+            .is_none()
+    }
 }
 
 #[cfg(feature = "iterator")]
@@ -1511,4 +1541,40 @@ mod test {
         assert_eq!(include.len(), 1);
         assert_eq!(include, vec![456]);
     }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn clear_works() {
+        const TEST_MAP: Map<&str, u32> = Map::new("test_map");
+
+        let mut storage = MockStorage::new();
+        TEST_MAP.save(&mut storage, "key0", &0u32).unwrap();
+        TEST_MAP.save(&mut storage, "key1", &1u32).unwrap();
+        TEST_MAP.save(&mut storage, "key2", &2u32).unwrap();
+        TEST_MAP.save(&mut storage, "key3", &3u32).unwrap();
+        TEST_MAP.save(&mut storage, "key4", &4u32).unwrap();
+
+        TEST_MAP.clear(&mut storage);
+
+        assert!(!TEST_MAP.has(&storage, "key0"));
+        assert!(!TEST_MAP.has(&storage, "key1"));
+        assert!(!TEST_MAP.has(&storage, "key2"));
+        assert!(!TEST_MAP.has(&storage, "key3"));
+        assert!(!TEST_MAP.has(&storage, "key4"));
+    }
+
+    #[test]
+    #[cfg(feature = "iterator")]
+    fn is_empty_works() {
+        const TEST_MAP: Map<&str, u32> = Map::new("test_map");
+
+        let mut storage = MockStorage::new();
+
+        assert!(TEST_MAP.is_empty(&storage));
+
+        TEST_MAP.save(&mut storage, "key1", &1u32).unwrap();
+        TEST_MAP.save(&mut storage, "key2", &2u32).unwrap();
+
+        assert!(!TEST_MAP.is_empty(&storage));
+    }
 }