Add SpreadAllocate and PackedAllocate traits and impls (#978)

* improve docs of PackedLayout methods

* add SpreadAllocate trait

Note: We cannot add allocate_spread to SpreadLayout trait since certain types (such as Result<T, E>) that implement SpreadLayout cannot implement SpreadAllocate; OR certain types can trivially implement SpreadLayout but are constraint with respect to SpreadAllocate and vice versa.

* add PackedAllocate trait

* add utility methods for SpreadAllocate and PackedAllocate traits

* improve macro hygiene of some trait impl code dupe macros

* normalize some import statement

* apply rustfmt

* implement {Spread,Packed}Allocate traits for primitive and prelude types

* implement {Packed,Spread}Allocate for std collections

* implement {Packed,Spread}Allocate for array primitives

* improve Codec impl for ink_storage::Pack<T>

* add {Packed,Spread}Allocate impl for ink_storage::Pack<T>

* add {Packed,Spread}Allocate impl for ink_storage::Memory

* add missing #[inline] annotations to ink_storage::Memory<T>

* implement {Packed,Spread}Allocate for tuple types

* add {Spread,Packed}Allocate impl for lazy storage abstractions

* add some #[inline] annotations

* add SpreadAllocate impl for Lazy<T>

* add SpreadAllocate impl for StorageVec

* add SpreadAllocate impl for StorageStash

* derive Default for Header

* remove unused import

* implement SpreadAllocate for the rest of the ink! storage collections

crates/storage/src/collections/binary_heap/children_vec.rs

@@ -26,6 +26,7 @@ use crate::{
     traits::{
         KeyPtr,
         PackedLayout,
+        SpreadAllocate,
         SpreadLayout,
     },
     Lazy,
@@ -307,6 +308,18 @@ where
     }
 }
 
+impl<T> SpreadAllocate for ChildrenVec<T>
+where
+    T: SpreadLayout + Ord + PackedLayout,
+{
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            len: SpreadAllocate::allocate_spread(ptr),
+            children: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}
+
 /// An iterator over shared references to the elements of the `BinaryHeap`.
 #[derive(Debug, Clone, Copy)]
 pub struct Iter<'a, T>

crates/storage/src/collections/binary_heap/storage.rs

@@ -21,6 +21,7 @@ use super::{
 use crate::traits::{
     KeyPtr,
     PackedLayout,
+    SpreadAllocate,
     SpreadLayout,
 };
 
@@ -98,3 +99,14 @@ where
         SpreadLayout::clear_spread(&self.elements, ptr);
     }
 }
+
+impl<T> SpreadAllocate for BinaryHeap<T>
+where
+    T: PackedLayout + Ord,
+{
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            elements: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

crates/storage/src/collections/bitstash/storage.rs

@@ -27,6 +27,7 @@ use crate::{
         forward_push_packed,
         KeyPtr,
         PackedLayout,
+        SpreadAllocate,
         SpreadLayout,
     },
 };
@@ -101,3 +102,12 @@ impl SpreadLayout for BitStash {
         SpreadLayout::clear_spread(&self.free, ptr);
     }
 }
+
+impl SpreadAllocate for BitStash {
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            counts: SpreadAllocate::allocate_spread(ptr),
+            free: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

crates/storage/src/collections/bitvec/storage.rs

@@ -23,6 +23,7 @@ use crate::{
         forward_push_packed,
         KeyPtr,
         PackedLayout,
+        SpreadAllocate,
         SpreadLayout,
     },
     Pack,
@@ -98,3 +99,12 @@ impl SpreadLayout for StorageBitvec {
         SpreadLayout::clear_spread(&self.bits, ptr);
     }
 }
+
+impl SpreadAllocate for StorageBitvec {
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            len: SpreadAllocate::allocate_spread(ptr),
+            bits: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

crates/storage/src/collections/hashmap/storage.rs

@@ -26,6 +26,7 @@ use crate::{
         forward_push_packed,
         KeyPtr,
         PackedLayout,
+        SpreadAllocate,
         SpreadLayout,
     },
 };
@@ -137,3 +138,18 @@ where
         SpreadLayout::clear_spread(&self.values, ptr);
     }
 }
+
+impl<K, V, H> SpreadAllocate for StorageHashMap<K, V, H>
+where
+    K: Ord + Clone + PackedLayout,
+    V: PackedLayout,
+    H: CryptoHash,
+    Key: From<<H as HashOutput>::Type>,
+{
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            keys: SpreadAllocate::allocate_spread(ptr),
+            values: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

crates/storage/src/collections/smallvec/storage.rs

@@ -16,6 +16,7 @@ use super::SmallVec;
 use crate::traits::{
     KeyPtr,
     PackedLayout,
+    SpreadAllocate,
     SpreadLayout,
 };
 
@@ -72,3 +73,15 @@ where
         SpreadLayout::clear_spread(&self.elems, ptr);
     }
 }
+
+impl<T, const N: usize> SpreadAllocate for SmallVec<T, N>
+where
+    T: PackedLayout,
+{
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            len: SpreadAllocate::allocate_spread(ptr),
+            elems: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

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

@@ -53,7 +53,7 @@ where
 }
 
 /// Stores general commonly required information about the storage stash.
-#[derive(Debug, scale::Encode, scale::Decode)]
+#[derive(Debug, Default, scale::Encode, scale::Decode)]
 #[cfg_attr(feature = "std", derive(scale_info::TypeInfo))]
 struct Header {
     /// The latest vacant index.

crates/storage/src/collections/stash/storage.rs

@@ -22,11 +22,14 @@ use super::{
 use crate::{
     lazy::LazyIndexMap,
     traits::{
+        forward_allocate_packed,
         forward_clear_packed,
         forward_pull_packed,
         forward_push_packed,
         KeyPtr,
+        PackedAllocate,
         PackedLayout,
+        SpreadAllocate,
         SpreadLayout,
     },
 };
@@ -88,12 +91,26 @@ impl SpreadLayout for Header {
     }
 }
 
+impl SpreadAllocate for Header {
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        forward_allocate_packed::<Self>(ptr)
+    }
+}
+
 impl PackedLayout for Header {
+    #[inline(always)]
     fn pull_packed(&mut self, _at: &Key) {}
+    #[inline(always)]
     fn push_packed(&self, _at: &Key) {}
+    #[inline(always)]
     fn clear_packed(&self, _at: &Key) {}
 }
 
+impl PackedAllocate for Header {
+    #[inline(always)]
+    fn allocate_packed(&mut self, _at: &Key) {}
+}
+
 impl<T> SpreadLayout for Entry<T>
 where
     T: PackedLayout,
@@ -161,3 +178,15 @@ where
         SpreadLayout::clear_spread(&self.entries, ptr);
     }
 }
+
+impl<T> SpreadAllocate for StorageStash<T>
+where
+    T: PackedLayout,
+{
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            header: SpreadAllocate::allocate_spread(ptr),
+            entries: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

crates/storage/src/collections/vec/storage.rs

@@ -20,6 +20,7 @@ use crate::{
     traits::{
         KeyPtr,
         PackedLayout,
+        SpreadAllocate,
         SpreadLayout,
     },
 };
@@ -77,3 +78,15 @@ where
         SpreadLayout::clear_spread(&self.elems, ptr);
     }
 }
+
+impl<T> SpreadAllocate for StorageVec<T>
+where
+    T: PackedLayout,
+{
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self {
+            len: SpreadAllocate::allocate_spread(ptr),
+            elems: SpreadAllocate::allocate_spread(ptr),
+        }
+    }
+}

crates/storage/src/lazy/cache_cell.rs

@@ -31,6 +31,7 @@ pub struct CacheCell<T: ?Sized> {
 
 impl<T> CacheCell<T> {
     /// Creates a new cache cell from the given value.
+    #[inline]
     pub fn new(value: T) -> Self {
         Self {
             inner: UnsafeCell::new(value),
@@ -54,6 +55,7 @@ where
 }
 
 impl<T> From<T> for CacheCell<T> {
+    #[inline]
     fn from(value: T) -> Self {
         Self::new(value)
     }
@@ -63,6 +65,7 @@ impl<T> Default for CacheCell<T>
 where
     T: Default,
 {
+    #[inline]
     fn default() -> Self {
         Self::new(<T as Default>::default())
     }
@@ -73,6 +76,7 @@ where
     T: ?Sized,
 {
     /// Returns a shared reference to the inner value.
+    #[inline]
     pub fn as_inner(&self) -> &T {
         // SAFETY: This is safe since we are returning a shared reference back
         //         to the caller while this method itself accesses `self` as
@@ -81,6 +85,7 @@ where
     }
 
     /// Returns an exclusive reference to the inner value.
+    #[inline]
     pub fn as_inner_mut(&mut self) -> &mut T {
         // SAFETY: This is safe since we are returning the exclusive reference
         //         of the inner value through the `get_mut` API which itself
@@ -90,6 +95,7 @@ where
     }
 
     /// Returns a mutable pointer to the inner value.
+    #[inline]
     pub fn get_ptr(&self) -> NonNull<T> {
         // SAFETY: The inner `T` of the internal `UnsafeCell` exists and thus
         //         the pointer that we get returned to it via `UnsafeCell::get`

crates/storage/src/lazy/entry.rs

@@ -21,7 +21,9 @@ use crate::traits::{
     push_spread_root_opt,
     ExtKeyPtr,
     KeyPtr,
+    PackedAllocate,
     PackedLayout,
+    SpreadAllocate,
     SpreadLayout,
 };
 use core::{
@@ -100,22 +102,36 @@ where
 {
     const FOOTPRINT: u64 = <T as SpreadLayout>::FOOTPRINT;
 
+    #[inline]
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
         let root_key = ExtKeyPtr::next_for::<Self>(ptr);
         Self::pull_spread_root(root_key)
     }
 
+    #[inline]
     fn push_spread(&self, ptr: &mut KeyPtr) {
         let root_key = ExtKeyPtr::next_for::<Self>(ptr);
         self.push_spread_root(root_key)
     }
 
+    #[inline]
     fn clear_spread(&self, ptr: &mut KeyPtr) {
         let root_key = ExtKeyPtr::next_for::<Self>(ptr);
         self.clear_spread_root(root_key)
     }
 }
 
+impl<T> SpreadAllocate for StorageEntry<T>
+where
+    T: SpreadLayout,
+{
+    #[inline]
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
+        Self::pull_spread_root(root_key)
+    }
+}
+
 impl<T> scale::Encode for StorageEntry<T>
 where
     T: scale::Encode,
@@ -173,6 +189,16 @@ where
     }
 }
 
+impl<T> PackedAllocate for StorageEntry<T>
+where
+    T: PackedAllocate,
+{
+    #[inline]
+    fn allocate_packed(&mut self, at: &Key) {
+        PackedAllocate::allocate_packed(&mut self.value, at)
+    }
+}
+
 impl<T> StorageEntry<T>
 where
     T: SpreadLayout,

crates/storage/src/lazy/lazy_array.rs

@@ -23,6 +23,7 @@ use crate::traits::{
     ExtKeyPtr,
     KeyPtr,
     PackedLayout,
+    SpreadAllocate,
     SpreadLayout,
 };
 use core::{
@@ -375,6 +376,7 @@ where
 {
     const FOOTPRINT: u64 = N as u64;
 
+    #[inline]
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
         Self::lazy(*ExtKeyPtr::next_for::<Self>(ptr))
     }
@@ -399,6 +401,16 @@ where
     }
 }
 
+impl<T, const N: usize> SpreadAllocate for LazyArray<T, N>
+where
+    T: PackedLayout,
+{
+    #[inline]
+    fn allocate_spread(ptr: &mut KeyPtr) -> Self {
+        Self::lazy(*ExtKeyPtr::next_for::<Self>(ptr))
+    }
+}
+
 impl<T, const N: usize> LazyArray<T, N> {
     /// Returns the offset key for the given index if not out of bounds.
     pub fn key_at(&self, at: Index) -> Option<Key> {