Add Mapping::contains(key) and `Mapping::insert_return_size(key, va…

…l)` (#1224)

* added storage_contains()

* SENTINEL const moved to on_chain env

* added Mapping::contains()

* added Mapping::insert_checked()

* CI satisfaction (then_some() -> then())

* CI satisfaction 2

* udpated env to return the size of prev stored val

* Mapping::insert_return_size()

* impl From<ReturnCode> for Option<u32>

* Apply suggestions from code review

Co-authored-by: Michael Müller <mich@elmueller.net>

Co-authored-by: Michael Müller <mich@elmueller.net>

crates/engine/src/ext.rs

@@ -228,20 +228,18 @@ impl Engine {
     }
 
     /// Writes the encoded value into the storage at the given key.
-    pub fn set_storage(&mut self, key: &[u8; 32], encoded_value: &[u8]) {
+    /// Returns the size of the previously stored value at the key if any.
+    pub fn set_storage(&mut self, key: &[u8; 32], encoded_value: &[u8]) -> Option<u32> {
         let callee = self.get_callee();
         let account_id = AccountId::from_bytes(&callee[..]);
 
         self.debug_info.inc_writes(account_id.clone());
         self.debug_info
             .record_cell_for_account(account_id, key.to_vec());
 
-        // We ignore if storage is already set for this key
-        let _ = self.database.insert_into_contract_storage(
-            &callee,
-            key,
-            encoded_value.to_vec(),
-        );
+        self.database
+            .insert_into_contract_storage(&callee, key, encoded_value.to_vec())
+            .map(|v| <u32>::try_from(v.len()).expect("usize to u32 conversion failed"))
     }
 
     /// Returns the decoded contract storage at the key if any.

crates/env/src/api.rs

@@ -183,12 +183,13 @@ where
     })
 }
 
-/// Writes the value to the contract storage under the given key.
+/// Writes the value to the contract storage under the given key and returns
+/// the size of pre-existing value at the specified key if any.
 ///
 /// # Panics
 ///
 /// - If the encode length of value exceeds the configured maximum value length of a storage entry.
-pub fn set_contract_storage<V>(key: &Key, value: &V)
+pub fn set_contract_storage<V>(key: &Key, value: &V) -> Option<u32>
 where
     V: scale::Encode,
 {
@@ -211,6 +212,16 @@ where
     })
 }
 
+/// Checks whether there is a value stored under the given key in
+/// the contract's storage.
+///
+/// If a value is stored under the specified key, the size of the value is returned.
+pub fn contract_storage_contains(key: &Key) -> Option<u32> {
+    <EnvInstance as OnInstance>::on_instance(|instance| {
+        EnvBackend::contract_storage_contains(instance, key)
+    })
+}
+
 /// Clears the contract's storage key entry.
 pub fn clear_contract_storage(key: &Key) {
     <EnvInstance as OnInstance>::on_instance(|instance| {

crates/env/src/backend.rs

@@ -162,8 +162,9 @@ impl CallFlags {
 
 /// Environmental contract functionality that does not require `Environment`.
 pub trait EnvBackend {
-    /// Writes the value to the contract storage under the given key.
-    fn set_contract_storage<V>(&mut self, key: &Key, value: &V)
+    /// Writes the value to the contract storage under the given key and returns
+    /// the size of the pre-existing value at the specified key if any.
+    fn set_contract_storage<V>(&mut self, key: &Key, value: &V) -> Option<u32>
     where
         V: scale::Encode;
 
@@ -176,6 +177,9 @@ pub trait EnvBackend {
     where
         R: scale::Decode;
 
+    /// Returns the size of a value stored under the specified key is returned if any.
+    fn contract_storage_contains(&mut self, key: &Key) -> Option<u32>;
+
     /// Clears the contract's storage key entry.
     fn clear_contract_storage(&mut self, key: &Key);
 

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

@@ -184,12 +184,12 @@ impl EnvInstance {
 }
 
 impl EnvBackend for EnvInstance {
-    fn set_contract_storage<V>(&mut self, key: &Key, value: &V)
+    fn set_contract_storage<V>(&mut self, key: &Key, value: &V) -> Option<u32>
     where
         V: scale::Encode,
     {
         let v = scale::Encode::encode(value);
-        self.engine.set_storage(key.as_ref(), &v[..]);
+        self.engine.set_storage(key.as_ref(), &v[..])
     }
 
     fn get_contract_storage<R>(&mut self, key: &Key) -> Result<Option<R>>
@@ -206,6 +206,12 @@ impl EnvBackend for EnvInstance {
         Ok(Some(decoded))
     }
 
+    fn contract_storage_contains(&mut self, _key: &Key) -> Option<u32> {
+        unimplemented!(
+            "the off-chain env does not implement `seal_contains_storage`, yet"
+        )
+    }
+
     fn clear_contract_storage(&mut self, key: &Key) {
         self.engine.clear_storage(key.as_ref())
     }

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

@@ -177,10 +177,24 @@ where
     }
 }
 
+/// Used as a sentinel value when reading and writing contract memory.
+///
+/// We use this value to signal `None` to a contract when only a primitive is allowed
+/// and we don't want to go through encoding a full Rust type. Using `u32::Max` is a safe
+/// sentinel because contracts are never allowed to use such a large amount of resources.
+/// So this value doesn't make sense for a memory location or length.
+const SENTINEL: u32 = u32::MAX;
+
 /// The raw return code returned by the host side.
 #[repr(transparent)]
 pub struct ReturnCode(u32);
 
+impl From<ReturnCode> for Option<u32> {
+    fn from(code: ReturnCode) -> Self {
+        (code.0 < SENTINEL).then(|| code.0)
+    }
+}
+
 impl ReturnCode {
     /// Returns the raw underlying `u32` representation.
     pub fn into_u32(self) -> u32 {
@@ -217,16 +231,12 @@ mod sys {
             data_len: u32,
         );
 
-        pub fn seal_set_storage(
-            key_ptr: Ptr32<[u8]>,
-            value_ptr: Ptr32<[u8]>,
-            value_len: u32,
-        );
         pub fn seal_get_storage(
             key_ptr: Ptr32<[u8]>,
             output_ptr: Ptr32Mut<[u8]>,
             output_len_ptr: Ptr32Mut<u32>,
         ) -> ReturnCode;
+
         pub fn seal_clear_storage(key_ptr: Ptr32<[u8]>);
 
         pub fn seal_call_chain_extension(
@@ -360,6 +370,14 @@ mod sys {
             output_ptr: Ptr32Mut<[u8]>,
             output_len_ptr: Ptr32Mut<u32>,
         );
+
+        pub fn seal_contains_storage(key_ptr: Ptr32<[u8]>) -> ReturnCode;
+
+        pub fn seal_set_storage(
+            key_ptr: Ptr32<[u8]>,
+            value_ptr: Ptr32<[u8]>,
+            value_len: u32,
+        ) -> ReturnCode;
     }
 }
 
@@ -475,14 +493,15 @@ pub fn deposit_event(topics: &[u8], data: &[u8]) {
     }
 }
 
-pub fn set_storage(key: &[u8], encoded_value: &[u8]) {
-    unsafe {
+pub fn set_storage(key: &[u8], encoded_value: &[u8]) -> Option<u32> {
+    let ret_code = unsafe {
         sys::seal_set_storage(
             Ptr32::from_slice(key),
             Ptr32::from_slice(encoded_value),
             encoded_value.len() as u32,
         )
-    }
+    };
+    ret_code.into()
 }
 
 pub fn clear_storage(key: &[u8]) {
@@ -504,6 +523,11 @@ pub fn get_storage(key: &[u8], output: &mut &mut [u8]) -> Result {
     ret_code.into()
 }
 
+pub fn storage_contains(key: &[u8]) -> Option<u32> {
+    let ret_code = unsafe { sys::seal_contains_storage(Ptr32::from_slice(key)) };
+    ret_code.into()
+}
+
 pub fn terminate(beneficiary: &[u8]) -> ! {
     unsafe { sys::seal_terminate(Ptr32::from_slice(beneficiary)) }
 }

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

@@ -219,12 +219,12 @@ impl EnvInstance {
 }
 
 impl EnvBackend for EnvInstance {
-    fn set_contract_storage<V>(&mut self, key: &Key, value: &V)
+    fn set_contract_storage<V>(&mut self, key: &Key, value: &V) -> Option<u32>
     where
         V: scale::Encode,
     {
         let buffer = self.scoped_buffer().take_encoded(value);
-        ext::set_storage(key.as_ref(), buffer);
+        ext::set_storage(key.as_ref(), buffer)
     }
 
     fn get_contract_storage<R>(&mut self, key: &Key) -> Result<Option<R>>
@@ -241,6 +241,10 @@ impl EnvBackend for EnvInstance {
         Ok(Some(decoded))
     }
 
+    fn contract_storage_contains(&mut self, key: &Key) -> Option<u32> {
+        ext::storage_contains(key.as_ref())
+    }
+
     fn clear_contract_storage(&mut self, key: &Key) {
         ext::clear_storage(key.as_ref())
     }

crates/storage/src/lazy/mapping.rs

@@ -134,6 +134,18 @@ where
         push_packed_root(value, &self.storage_key(&key));
     }
 
+    /// Insert the given `value` to the contract storage.
+    ///
+    /// Returns the size of the pre-existing value at the specified key if any.
+    #[inline]
+    pub fn insert_return_size<Q, R>(&mut self, key: Q, value: &R) -> Option<u32>
+    where
+        Q: scale::EncodeLike<K>,
+        R: scale::EncodeLike<V> + PackedLayout,
+    {
+        push_packed_root(value, &self.storage_key(&key))
+    }
+
     /// Get the `value` at `key` from the contract storage.
     ///
     /// Returns `None` if no `value` exists at the given `key`.
@@ -145,6 +157,17 @@ where
         pull_packed_root_opt(&self.storage_key(&key))
     }
 
+    /// Get the size of a value stored at `key` in the contract storage.
+    ///
+    /// Returns `None` if no `value` exists at the given `key`.
+    #[inline]
+    pub fn contains<Q>(&self, key: Q) -> Option<u32>
+    where
+        Q: scale::EncodeLike<K>,
+    {
+        ink_env::contract_storage_contains(&self.storage_key(&key))
+    }
+
     /// Clears the value at `key` from storage.
     pub fn remove<Q>(&self, key: Q)
     where

crates/storage/src/traits/impls/mod.rs

@@ -162,7 +162,7 @@ pub fn forward_push_packed<T>(entity: &T, ptr: &mut KeyPtr)
 where
     T: PackedLayout,
 {
-    push_packed_root::<T>(entity, ptr.next_for::<T>())
+    push_packed_root::<T>(entity, ptr.next_for::<T>());
 }
 
 /// Clears an instance of type `T` in packed fashion from the contract storage.

crates/storage/src/traits/mod.rs

@@ -202,12 +202,12 @@ where
 ///   packed layout.
 /// - Users should prefer using this function directly instead of using the
 ///   trait methods on [`PackedLayout`].
-pub fn push_packed_root<T>(entity: &T, root_key: &Key)
+pub fn push_packed_root<T>(entity: &T, root_key: &Key) -> Option<u32>
 where
     T: PackedLayout,
 {
     <T as PackedLayout>::push_packed(entity, root_key);
-    ink_env::set_contract_storage(root_key, entity);
+    ink_env::set_contract_storage(root_key, entity)
 }
 
 /// Clears the entity from the contract storage using packed layout.

crates/storage/src/traits/optspec.rs

@@ -127,7 +127,7 @@ where
             //
             // Sadly this does not work well with `Option<Option<T>>`.
             // For this we'd need specialization in Rust or similar.
-            super::push_packed_root(value, root_key)
+            super::push_packed_root(value, root_key);
         }
         None => {
             // Clear the associated storage cell since the entity is `None`.