diff --git a/crates/storage/src/lazy/mapping.rs b/crates/storage/src/lazy/mapping.rs index f5b9fb8aab7..6bc9c2f271a 100644 --- a/crates/storage/src/lazy/mapping.rs +++ b/crates/storage/src/lazy/mapping.rs @@ -74,7 +74,7 @@ where Q: scale::EncodeLike, R: scale::EncodeLike + PackedLayout, { - push_packed_root(value, &self.storage_key(key)); + push_packed_root(value, &self.storage_key(&key)); } /// Get the `value` at `key` from the contract storage. @@ -85,14 +85,30 @@ where where Q: scale::EncodeLike, { - pull_packed_root_opt(&self.storage_key(key)) + pull_packed_root_opt(&self.storage_key(&key)) + } + + /// Clears the value at `key` from storage. + pub fn remove(&self, key: Q) + where + Q: scale::EncodeLike, + { + let storage_key = self.storage_key(&key); + if ::REQUIRES_DEEP_CLEAN_UP { + // There are types which need to perform some action before being cleared. Here we + // indicate to those types that they should start tidying up. + if let Some(value) = self.get(key) { + ::clear_packed(&value, &storage_key); + } + } + ink_env::clear_contract_storage(&storage_key); } /// Returns a `Key` pointer used internally by the storage API. /// /// This key is a combination of the `Mapping`'s internal `offset_key` /// and the user provided `key`. - fn storage_key(&self, key: Q) -> Key + fn storage_key(&self, key: &Q) -> Key where Q: scale::EncodeLike, { @@ -186,4 +202,56 @@ mod tests { }) .unwrap() } + + #[test] + fn can_clear_entries() { + ink_env::test::run_test::(|_| { + // We use `Pack` here since it `REQUIRES_DEEP_CLEAN_UP` + use crate::Pack; + + // Given + let mut mapping: Mapping = Mapping::new([0u8; 32].into()); + let mut deep_mapping: Mapping> = Mapping::new([1u8; 32].into()); + + mapping.insert(&1, &2); + assert_eq!(mapping.get(&1), Some(2)); + + deep_mapping.insert(&1u8, &Pack::new(Pack::new(2u8))); + assert_eq!(deep_mapping.get(&1), Some(Pack::new(2u8))); + + // When + mapping.remove(&1); + deep_mapping.remove(&1); + + // Then + assert_eq!(mapping.get(&1), None); + assert_eq!(deep_mapping.get(&1), None); + + Ok(()) + }) + .unwrap() + } + + #[test] + fn can_clear_unexistent_entries() { + ink_env::test::run_test::(|_| { + // We use `Pack` here since it `REQUIRES_DEEP_CLEAN_UP` + use crate::Pack; + + // Given + let mapping: Mapping = Mapping::new([0u8; 32].into()); + let deep_mapping: Mapping> = Mapping::new([1u8; 32].into()); + + // When + mapping.remove(&1); + deep_mapping.remove(&1); + + // Then + assert_eq!(mapping.get(&1), None); + assert_eq!(deep_mapping.get(&1), None); + + Ok(()) + }) + .unwrap() + } }