[storage] Implement Drop for Pack

crates/primitives/src/key_ptr.rs

@@ -43,7 +43,7 @@ impl KeyPtr {
     }
 
     /// Returns the underlying offset key.
-    pub fn key(&self) -> Key {
-        self.key
+    pub fn key(&self) -> &Key {
+        &self.key
     }
 }

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

@@ -33,7 +33,6 @@ use crate::{
 };
 
 pub use children_vec::Iter;
-use ink_primitives::Key;
 pub use reverse::Reverse;
 
 /// A priority queue implemented with a binary heap.
@@ -51,14 +50,6 @@ where
 {
     /// The individual elements of the heap.
     elements: ChildrenVec<T>,
-    /// The offset key for the N cells.
-    ///
-    /// If the lazy chunk has been initialized during contract initialization
-    /// the key will be `None` since there won't be a storage region associated
-    /// to the lazy chunk which prevents it from lazily loading elements. This,
-    /// however, is only checked at contract runtime. We might incorporate
-    /// compile-time checks for this particular use case later on.
-    key: Option<Key>,
 }
 
 impl<T> BinaryHeap<T>
@@ -69,7 +60,6 @@ where
     pub fn new() -> Self {
         Self {
             elements: ChildrenVec::new(),
-            key: None,
         }
     }
 

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

@@ -85,10 +85,8 @@ where
     const FOOTPRINT: u64 = <ChildrenVec<T> as SpreadLayout>::FOOTPRINT;
 
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
-        let key = ptr.key();
         Self {
             elements: SpreadLayout::pull_spread(ptr),
-            key: Some(key),
         }
     }
 

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

@@ -210,10 +210,7 @@ fn spread_layout_push_pull_works() -> ink_env::Result<()> {
         // both instances are equal:
         let heap2 =
             <BinaryHeap<u8> as SpreadLayout>::pull_spread(&mut KeyPtr::from(root_key));
-        // we compare only the `elements` of the heap, since for `heap1` the `heap1.key`
-        // is `None`, since no pull was executed on this object yet. `heap2.key` will be
-        // `Some(root_key)` though, since we just pulled it from storage.
-        assert_eq!(heap1.elements, heap2.elements);
+        assert_eq!(heap1, heap2);
         Ok(())
     })
 }

crates/storage/src/pack.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use crate::traits::{
-    clear_spread_root_opt,
+    clear_spread_root,
     forward_clear_packed,
     forward_pull_packed,
     forward_push_packed,
@@ -117,6 +117,14 @@ where
         }
     }
 
+    /// Creates a new packed value.
+    pub fn new_with_key(value: T, key: Key) -> Self {
+        Self {
+            inner: value,
+            key: Some(key),
+        }
+    }
+
     /// Returns a shared reference to the packed value.
     pub fn as_inner(pack: &Pack<T>) -> &T {
         &pack.inner
@@ -126,12 +134,6 @@ where
     pub fn as_inner_mut(pack: &mut Pack<T>) -> &mut T {
         &mut pack.inner
     }
-
-    /// Sets the key.
-    fn set_key(mut self, value: Option<Key>) -> Self {
-        self.key = value;
-        self
-    }
 }
 
 impl<T> Drop for Pack<T>
@@ -140,7 +142,7 @@ where
 {
     fn drop(&mut self) {
         if let Some(key) = self.key {
-            clear_spread_root_opt::<T, _>(&key, || Some(&self.inner))
+            clear_spread_root::<T>(&self.inner, &key)
         }
     }
 }
@@ -172,7 +174,8 @@ where
     const FOOTPRINT: u64 = 1;
 
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
-        Pack::from(forward_pull_packed::<T>(ptr)).set_key(Some(ptr.key()))
+        let inner = forward_pull_packed::<T>(ptr);
+        Pack::new_with_key(inner, *ptr.key())
     }
 
     fn push_spread(&self, ptr: &mut KeyPtr) {