Add Mapping::take()

CHANGELOG.md

@@ -6,6 +6,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 - Allows to use `Result<Self, Error>` as a return type in constructors - [#1446](https://github.com/paritytech/ink/pull/1446)
+- Add `Mapping::take()` function allowing to get a value removing it from storage - [#1461](https://github.com/paritytech/ink/pull/1461)
 - Introduces conditional compilation to messages, constructors and events - [#1458](https://github.com/paritytech/ink/pull/1458)
 
 - Remove random function from ink!.

crates/engine/src/ext.rs

@@ -253,6 +253,22 @@ impl Engine {
         }
     }
 
+    /// Removes the storage entries at the given key,
+    /// returning previously stored value at the key if any.
+    pub fn take_storage(&mut self, key: &[u8], output: &mut &mut [u8]) -> Result {
+        let callee = self.get_callee();
+        let account_id = AccountId::from_bytes(&callee[..]);
+
+        self.debug_info.inc_writes(account_id);
+        match self.database.remove_contract_storage(&callee, key) {
+            Some(val) => {
+                set_output(output, &val);
+                Ok(())
+            }
+            None => Err(Error::KeyNotFound),
+        }
+    }
+
     /// Returns the size of the value stored in the contract storage at the key if any.
     pub fn contains_storage(&mut self, key: &[u8]) -> Option<u32> {
         let callee = self.get_callee();

crates/env/src/api.rs

@@ -214,6 +214,21 @@ where
     })
 }
 
+/// Removes the `value` at `key`,
+/// returning the previous `value` at `key` from storage
+/// # Errors
+///
+/// - If the decoding of the typed value failed (`KeyNotFound`)
+pub fn take_contract_storage<K, R>(key: &K) -> Result<Option<R>>
+where
+    K: scale::Encode,
+    R: Storable,
+{
+    <EnvInstance as OnInstance>::on_instance(|instance| {
+        EnvBackend::take_contract_storage::<K, R>(instance, key)
+    })
+}
+
 /// Checks whether there is a value stored under the given storage key in the contract's storage.
 ///
 /// If a value is stored under the specified key, the size of the value is returned.

crates/env/src/backend.rs

@@ -180,6 +180,17 @@ pub trait EnvBackend {
         K: scale::Encode,
         R: Storable;
 
+    /// Removes the `value` at `key`,
+    /// returning the previous `value` at `key` from storage
+    /// if any
+    /// # Errors
+    ///
+    /// - If the decoding of the typed value failed
+    fn take_contract_storage<K, R>(&mut self, key: &K) -> Result<Option<R>>
+    where
+        K: scale::Encode,
+        R: Storable;
+
     /// Returns the size of a value stored under the given storage key is returned if any.
     fn contains_contract_storage<K>(&mut self, key: &K) -> Option<u32>
     where

crates/env/src/engine/off_chain/impls.rs

@@ -209,6 +209,24 @@ impl EnvBackend for EnvInstance {
         Ok(Some(decoded))
     }
 
+    fn take_contract_storage<K, R>(&mut self, key: &K) -> Result<Option<R>>
+    where
+        K: scale::Encode,
+        R: Storable,
+    {
+        let mut output: [u8; 9600] = [0; 9600];
+        match self
+            .engine
+            .take_storage(&key.encode(), &mut &mut output[..])
+        {
+            Ok(_) => (),
+            Err(ext::Error::KeyNotFound) => return Ok(None),
+            Err(_) => panic!("encountered unexpected error"),
+        }
+        let decoded = Storable::decode(&mut &output[..])?;
+        Ok(Some(decoded))
+    }
+
     fn contains_contract_storage<K>(&mut self, key: &K) -> Option<u32>
     where
         K: scale::Encode,

crates/env/src/engine/on_chain/ext.rs

@@ -418,6 +418,26 @@ mod sys {
             out_ptr: Ptr32Mut<[u8]>,
             out_len_ptr: Ptr32Mut<u32>,
         ) -> ReturnCode;
+
+        /// Retrieve and remove the value under the given key from storage.
+        ///
+        /// # Parameters
+        ///
+        /// - `key_ptr`: pointer into the linear memory where the key of the requested value is placed.
+        /// - `key_len`: the length of the key in bytes.
+        /// - `out_ptr`: pointer to the linear memory where the value is written to.
+        /// - `out_len_ptr`: in-out pointer into linear memory where the buffer length is read from and
+        ///   the value length is written to.
+        ///
+        /// # Errors
+        ///
+        /// `ReturnCode::KeyNotFound`
+        pub fn seal_take_storage(
+            key_ptr: Ptr32<[u8]>,
+            key_len: u32,
+            out_ptr: Ptr32Mut<[u8]>,
+            out_len_ptr: Ptr32Mut<u32>,
+        ) -> ReturnCode;
     }
 }
 
@@ -572,6 +592,23 @@ pub fn get_storage(key: &[u8], output: &mut &mut [u8]) -> Result {
     ret_code.into()
 }
 
+#[inline(always)]
+pub fn take_storage(key: &[u8], output: &mut &mut [u8]) -> Result {
+    let mut output_len = output.len() as u32;
+    let ret_code = {
+        unsafe {
+            sys::seal_take_storage(
+                Ptr32::from_slice(key),
+                key.len() as u32,
+                Ptr32Mut::from_slice(output),
+                Ptr32Mut::from_ref(&mut output_len),
+            )
+        }
+    };
+    extract_from_slice(output, output_len as usize);
+    ret_code.into()
+}
+
 pub fn storage_contains(key: &[u8]) -> Option<u32> {
     let ret_code =
         unsafe { sys::seal_contains_storage(Ptr32::from_slice(key), key.len() as u32) };

crates/env/src/engine/on_chain/impls.rs

@@ -251,6 +251,23 @@ impl EnvBackend for EnvInstance {
         Ok(Some(decoded))
     }
 
+    fn take_contract_storage<K, R>(&mut self, key: &K) -> Result<Option<R>>
+    where
+        K: scale::Encode,
+        R: Storable,
+    {
+        let mut buffer = self.scoped_buffer();
+        let key = buffer.take_encoded(key);
+        let output = &mut buffer.take_rest();
+        match ext::take_storage(key, output) {
+            Ok(_) => (),
+            Err(ExtError::KeyNotFound) => return Ok(None),
+            Err(_) => panic!("encountered unexpected error"),
+        }
+        let decoded = Storable::decode(&mut &output[..])?;
+        Ok(Some(decoded))
+    }
+
     fn contains_contract_storage<K>(&mut self, key: &K) -> Option<u32>
     where
         K: scale::Encode,

crates/storage/src/lazy/mapping.rs

@@ -164,6 +164,19 @@ where
             .unwrap_or_else(|error| panic!("Failed to get value in Mapping: {:?}", error))
     }
 
+    /// Removes the `value` at `key`, returning the previous `value` at `key` from storage.
+    ///
+    /// Returns `None` if no `value` exists at the given `key`.
+    #[inline]
+    pub fn take<Q>(&self, key: Q) -> Option<V>
+    where
+        Q: scale::EncodeLike<K>,
+    {
+        ink_env::take_contract_storage(&(&KeyType::KEY, key)).unwrap_or_else(|error| {
+            panic!("Failed to take value in Mapping: {:?}", error)
+        })
+    }
+
     /// Get the size of a value stored at `key` in the contract storage.
     ///
     /// Returns `None` if no `value` exists at the given `key`.
@@ -294,6 +307,30 @@ mod tests {
         .unwrap()
     }
 
+    #[test]
+    fn insert_and_take_work() {
+        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
+            let mut mapping: Mapping<u8, _> = Mapping::new();
+            mapping.insert(&1, &2);
+            assert_eq!(mapping.take(&1), Some(2));
+            assert!(mapping.get(&1).is_none());
+
+            Ok(())
+        })
+        .unwrap()
+    }
+
+    #[test]
+    fn take_empty_value_work() {
+        ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
+            let mapping: Mapping<u8, u8> = Mapping::new();
+            assert_eq!(mapping.take(&1), None);
+
+            Ok(())
+        })
+        .unwrap()
+    }
+
     #[test]
     fn can_clear_entries() {
         ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {