Simple Mapping type improvements (#979)

* Add `Mapping` storage collection

* Implement `insert` and `get` for `Mapping`

* Implement `SpreadLayout` for `Mapping`

* Fix typo

* Add some basic tests

* Fix some documentation formatting

* Use `PackedLayout` as trait bound instead of `Encode/Decode`

* Avoid using low level `ink_env` functions when interacting with storage

* RustFmt

* Appease Clippy

* Only use single `PhantomData` field

* Change `get` API to take reference to `key`

* Implement `TypeInfo` and `StorageLayout` for `Mapping`

* Properly gate `TypeInfo` and `StorageLayout` impls behind `std`

* Replace `HashMap` with `Mapping` in ERC-20 example

* Return `Option` from `Mapping::get`

* Update ERC-20 to handle `Option` returns

* Change `get` and `key` to use `Borrow`-ed values

* Add `Debug` and `Default` implementations

* Proper spelling

* Change `insert` to only accept borrowed K,V pairs

* Update ERC-20 example accordingly

* Make more explicit what each `key` is referring to

* Try using a `RefCell` instead of passing `Key` around

* Try using `UnsafeCell` instead

* Revert "Try using a `RefCell` instead of passing `Key` around"

This reverts commit cede033.

Using `RefCell`/`UnsafeCell` doesn't reduce the contract size more than
what we have now, and it introduced `unsafe` code. We believe the
limiting factor here is the `Key` type definition anyways.

* Clean up some of the documentation

* adjust the Mapping type for the new SpreadAllocate trait

* adjust ERC-20 example for changes in Mapping type

* remove commented out code

* add doc comment to new_init

* make it possible to use references in more cases with Mapping

* use references in more cases for ERC-20 example contract

* remove unnecessary references in Mapping methods

* refactor/improve pull_packed_root_opt utility method slightly

* fix ERC-20 example contract

The problem with *self.total_supply is that it may implicitly read from storage in case it has not yet read a value from storage whereas Lazy::set just writes the value to the Lazy instance.

Co-authored-by: Hernando Castano <[email protected]>
Co-authored-by: Hernando Castano <[email protected]>

Author: Robbepop

crates/storage/src/collections/mapping/mod.rs

@@ -14,21 +14,20 @@
 
 //! A simple mapping to contract storage.
 //!
-//! This mapping doesn't actually "own" any data. Instead it is just a simple wrapper around the
-//! contract storage facilities.
+//! # Note
+//!
+//! This mapping doesn't actually "own" any data.
+//! Instead it is just a simple wrapper around the contract storage facilities.
 
 use crate::traits::{
-    clear_spread_root,
     pull_packed_root_opt,
-    pull_spread_root,
     push_packed_root,
-    push_spread_root,
     ExtKeyPtr,
     KeyPtr,
     PackedLayout,
+    SpreadAllocate,
     SpreadLayout,
 };
-
 use core::marker::PhantomData;
 
 use ink_env::hash::{
@@ -42,7 +41,7 @@ use ink_primitives::Key;
 #[derive(Default)]
 pub struct Mapping<K, V> {
     offset_key: Key,
-    _marker: PhantomData<(K, V)>,
+    _marker: PhantomData<fn() -> (K, V)>,
 }
 
 impl<K, V> core::fmt::Debug for Mapping<K, V> {
@@ -53,52 +52,49 @@ impl<K, V> core::fmt::Debug for Mapping<K, V> {
     }
 }
 
-impl<K, V> Mapping<K, V>
-where
-    K: PackedLayout,
-    V: PackedLayout,
-{
+impl<K, V> Mapping<K, V> {
     /// Creates a new empty `Mapping`.
-    ///
-    /// TODO [#961]: Ideally we improve how this is initialized by extending the
-    /// `SpreadLayout`/`PackedLayout` traits for non-caching data structures.
-    pub fn new(offset_key: Key) -> Self {
+    fn new(offset_key: Key) -> Self {
         Self {
             offset_key,
             _marker: Default::default(),
         }
     }
+}
 
+impl<K, V> Mapping<K, V>
+where
+    K: PackedLayout,
+    V: PackedLayout,
+{
     /// Insert the given `value` to the contract storage.
-    pub fn insert<Q, R>(&mut self, key: &Q, value: &R)
+    #[inline]
+    pub fn insert<Q, R>(&mut self, key: Q, value: &R)
     where
-        K: core::borrow::Borrow<Q>,
-        Q: scale::Encode,
-        V: core::borrow::Borrow<R>,
-        R: PackedLayout,
+        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`.
-    pub fn get<Q>(&self, key: &Q) -> Option<V>
+    #[inline]
+    pub fn get<Q>(&self, key: Q) -> Option<V>
     where
-        K: core::borrow::Borrow<Q>,
-        Q: scale::Encode,
+        Q: scale::EncodeLike<K>,
     {
         pull_packed_root_opt(&self.storage_key(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<Q>(&self, key: &Q) -> Key
+    /// This key is a combination of the `Mapping`'s internal `offset_key`
+    /// and the user provided `key`.
+    fn storage_key<Q>(&self, key: Q) -> Key
     where
-        K: core::borrow::Borrow<Q>,
-        Q: scale::Encode,
+        Q: scale::EncodeLike<K>,
     {
         let encodedable_key = (&self.offset_key, key);
         let mut output = <Blake2x256 as HashOutput>::Type::default();
@@ -113,20 +109,32 @@ impl<K, V> SpreadLayout for Mapping<K, V> {
 
     #[inline]
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
-        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
-        pull_spread_root::<Self>(root_key)
+        // Note: There is no need to pull anything from the storage for the
+        //       mapping type since it initializes itself entirely by the
+        //       given key pointer.
+        Self::new(*ExtKeyPtr::next_for::<Self>(ptr))
     }
 
     #[inline]
     fn push_spread(&self, ptr: &mut KeyPtr) {
-        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
-        push_spread_root::<Self>(self, root_key);
+        // Note: The mapping type does not store any state in its associated
+        //       storage region, therefore only the pointer has to be incremented.
+        ptr.advance_by(Self::FOOTPRINT);
     }
 
     #[inline]
     fn clear_spread(&self, ptr: &mut KeyPtr) {
-        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
-        clear_spread_root::<Self>(self, root_key);
+        // Note: The mapping type is not aware of its elements, therefore
+        //       it is not possible to clean up after itself.
+        ptr.advance_by(Self::FOOTPRINT);
+    }
+}
+
+impl<K, V> SpreadAllocate for Mapping<K, V> {
+    #[inline]
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        // Note: The mapping type initializes itself entirely by the key pointer.
+        Self::new(*ExtKeyPtr::next_for::<Self>(ptr))
     }
 }
 

crates/storage/src/traits/optspec.rs

@@ -102,17 +102,19 @@ pub fn pull_packed_root_opt<T>(root_key: &Key) -> Option<T>
 where
     T: PackedLayout,
 {
-    match ink_env::get_contract_storage::<T>(root_key)
-        .expect("decoding does not match expected type")
-    {
-        Some(mut value) => {
-            // In case the contract storage is occupied we handle
-            // the Option<T> as if it was a T.
+    ink_env::get_contract_storage::<T>(root_key)
+        .unwrap_or_else(|error| {
+            panic!(
+                "failed to pull packed from root key {}: {:?}",
+                root_key, error
+            )
+        })
+        .map(|mut value| {
+            // In case the contract storage is occupied at the root key
+            // we handle the Option<T> as if it was a T.
             <T as PackedLayout>::pull_packed(&mut value, root_key);
-            Some(value)
-        }
-        None => None,
-    }
+            value
+        })
 }
 
 pub fn push_packed_root_opt<T>(entity: Option<&T>, root_key: &Key)
@@ -128,7 +130,7 @@ where
             super::push_packed_root(value, root_key)
         }
         None => {
-            // Clear the associated storage cell.
+            // Clear the associated storage cell since the entity is `None`.
             ink_env::clear_contract_storage(root_key);
         }
     }

examples/erc20/lib.rs

@@ -4,9 +4,14 @@ use ink_lang as ink;
 
 #[ink::contract]
 mod erc20 {
+    use ink_primitives::{
+        Key,
+        KeyPtr,
+    };
     use ink_storage::{
         collections::mapping::Mapping,
         lazy::Lazy,
+        traits::SpreadAllocate,
     };
 
     /// A simple ERC-20 contract.
@@ -55,24 +60,37 @@ mod erc20 {
     /// The ERC-20 result type.
     pub type Result<T> = core::result::Result<T, Error>;
 
+    impl SpreadAllocate for Erc20 {
+        fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+            Self {
+                total_supply: SpreadAllocate::allocate_spread(ptr),
+                balances: SpreadAllocate::allocate_spread(ptr),
+                allowances: SpreadAllocate::allocate_spread(ptr),
+            }
+        }
+    }
+
     impl Erc20 {
         /// Creates a new ERC-20 contract with the specified initial supply.
         #[ink(constructor)]
         pub fn new(initial_supply: Balance) -> Self {
+            let root_key = Key::from([0x00; 32]);
+            let mut key_ptr = KeyPtr::from(root_key);
+            let mut instance = Self::allocate_spread(&mut key_ptr);
+            instance.new_init(initial_supply);
+            instance
+        }
+
+        /// Default initializes the ERC-20 contract with the specified initial supply.
+        fn new_init(&mut self, initial_supply: Balance) {
             let caller = Self::env().caller();
-            let mut balances = Mapping::new([1u8; 32].into());
-            balances.insert(&caller, &initial_supply);
-            let instance = Self {
-                total_supply: Lazy::new(initial_supply),
-                balances,
-                allowances: Mapping::new([1u8; 32].into()),
-            };
+            self.balances.insert(&caller, &initial_supply);
+            Lazy::set(&mut self.total_supply, initial_supply);
             Self::env().emit_event(Transfer {
                 from: None,
                 to: Some(caller),
                 value: initial_supply,
             });
-            instance
         }
 
         /// Returns the total token supply.
@@ -86,15 +104,41 @@ mod erc20 {
         /// Returns `0` if the account is non-existent.
         #[ink(message)]
         pub fn balance_of(&self, owner: AccountId) -> Balance {
-            self.balances.get(&owner).unwrap_or_default() // .copied().unwrap_or(0)
+            self.balance_of_impl(&owner)
+        }
+
+        /// Returns the account balance for the specified `owner`.
+        ///
+        /// Returns `0` if the account is non-existent.
+        ///
+        /// # Note
+        ///
+        /// Prefer to call this method over `balance_of` since this
+        /// works using references which are more efficient in Wasm.
+        #[inline]
+        fn balance_of_impl(&self, owner: &AccountId) -> Balance {
+            self.balances.get(owner).unwrap_or_default()
         }
 
         /// Returns the amount which `spender` is still allowed to withdraw from `owner`.
         ///
         /// Returns `0` if no allowance has been set `0`.
         #[ink(message)]
         pub fn allowance(&self, owner: AccountId, spender: AccountId) -> Balance {
-            self.allowances.get(&(owner, spender)).unwrap_or_default() //.copied().unwrap_or(0)
+            self.allowance_impl(&owner, &spender)
+        }
+
+        /// Returns the amount which `spender` is still allowed to withdraw from `owner`.
+        ///
+        /// Returns `0` if no allowance has been set `0`.
+        ///
+        /// # Note
+        ///
+        /// Prefer to call this method over `allowance` since this
+        /// works using references which are more efficient in Wasm.
+        #[inline]
+        fn allowance_impl(&self, owner: &AccountId, spender: &AccountId) -> Balance {
+            self.allowances.get((owner, spender)).unwrap_or_default()
         }
 
         /// Transfers `value` amount of tokens from the caller's account to account `to`.
@@ -108,7 +152,7 @@ mod erc20 {
         #[ink(message)]
         pub fn transfer(&mut self, to: AccountId, value: Balance) -> Result<()> {
             let from = self.env().caller();
-            self.transfer_from_to(from, to, value)
+            self.transfer_from_to(&from, &to, value)
         }
 
         /// Allows `spender` to withdraw from the caller's account multiple times, up to
@@ -120,7 +164,7 @@ mod erc20 {
         #[ink(message)]
         pub fn approve(&mut self, spender: AccountId, value: Balance) -> Result<()> {
             let owner = self.env().caller();
-            self.allowances.insert(&(owner, spender), &value);
+            self.allowances.insert((&owner, &spender), &value);
             self.env().emit_event(Approval {
                 owner,
                 spender,
@@ -151,12 +195,13 @@ mod erc20 {
             value: Balance,
         ) -> Result<()> {
             let caller = self.env().caller();
-            let allowance = self.allowance(from, caller);
+            let allowance = self.allowance_impl(&from, &caller);
             if allowance < value {
                 return Err(Error::InsufficientAllowance)
             }
-            self.transfer_from_to(from, to, value)?;
-            self.allowances.insert(&(from, caller), &(allowance - value));
+            self.transfer_from_to(&from, &to, value)?;
+            self.allowances
+                .insert((&from, &caller), &(allowance - value));
             Ok(())
         }
 
@@ -170,20 +215,21 @@ mod erc20 {
         /// the caller's account balance.
         fn transfer_from_to(
             &mut self,
-            from: AccountId,
-            to: AccountId,
+            from: &AccountId,
+            to: &AccountId,
             value: Balance,
         ) -> Result<()> {
-            let from_balance = self.balance_of(from);
+            let from_balance = self.balance_of_impl(from);
             if from_balance < value {
                 return Err(Error::InsufficientBalance)
             }
-            self.balances.insert(&from, &(from_balance - value));
-            let to_balance = self.balance_of(to);
-            self.balances.insert(&to, &(to_balance + value));
+
+            self.balances.insert(from, &(from_balance - value));
+            let to_balance = self.balance_of_impl(to);
+            self.balances.insert(to, &(to_balance + value));
             self.env().emit_event(Transfer {
-                from: Some(from),
-                to: Some(to),
+                from: Some(*from),
+                to: Some(*to),
                 value,
             });
             Ok(())