[storage] Implement Drop for Pack

crates/primitives/src/key_ptr.rs

@@ -41,4 +41,9 @@ impl KeyPtr {
         self.key += old_shift;
         &self.key
     }
+
+    /// Returns the underlying offset key.
+    pub fn key(&self) -> Key {
+        self.key
+    }
 }

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

@@ -33,6 +33,7 @@ use crate::{
 };
 
 pub use children_vec::Iter;
+use ink_primitives::Key;
 pub use reverse::Reverse;
 
 /// A priority queue implemented with a binary heap.
@@ -50,6 +51,14 @@ 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>
@@ -60,6 +69,7 @@ where
     pub fn new() -> Self {
         Self {
             elements: ChildrenVec::new(),
+            key: None,
         }
     }
 

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

@@ -85,8 +85,10 @@ 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,7 +210,10 @@ 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));
-        assert_eq!(heap1, heap2);
+        // 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);
         Ok(())
     })
 }
@@ -255,6 +258,16 @@ fn drop_works() {
 
         assert!(setup_result.is_ok(), "setup should not panic");
 
+        let contract_id = ink_env::test::get_current_contract_account_id::<
+            ink_env::DefaultEnvironment,
+        >()
+        .expect("Cannot get contract id");
+        let used_cells = ink_env::test::count_used_storage_cells::<
+            ink_env::DefaultEnvironment,
+        >(&contract_id)
+        .expect("Used cells must be returned");
+        assert_eq!(used_cells, 0);
+
         let _ =
             <BinaryHeap<u8> as SpreadLayout>::pull_spread(&mut KeyPtr::from(root_key));
         Ok(())

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

@@ -352,3 +352,38 @@ fn storage_is_cleared_completely_after_pull_lazy() {
     })
     .unwrap()
 }
+
+#[test]
+#[should_panic(expected = "storage entry was empty")]
+fn drop_works() {
+    ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
+        let root_key = Key::from([0x42; 32]);
+
+        // if the setup panics it should not cause the test to pass
+        let setup_result = std::panic::catch_unwind(|| {
+            let hmap = filled_hmap();
+            SpreadLayout::push_spread(&hmap, &mut KeyPtr::from(root_key));
+            let _ = <StorageHashMap<u8, i32> as SpreadLayout>::pull_spread(
+                &mut KeyPtr::from(root_key),
+            );
+            // hmap is dropped which should clear the cells
+        });
+        assert!(setup_result.is_ok(), "setup should not panic");
+
+        let contract_id = ink_env::test::get_current_contract_account_id::<
+            ink_env::DefaultEnvironment,
+        >()
+        .expect("Cannot get contract id");
+        let used_cells = ink_env::test::count_used_storage_cells::<
+            ink_env::DefaultEnvironment,
+        >(&contract_id)
+        .expect("used cells must be returned");
+        assert_eq!(used_cells, 0);
+
+        let _ = <StorageHashMap<u8, i32> as SpreadLayout>::pull_spread(
+            &mut KeyPtr::from(root_key),
+        );
+        Ok(())
+    })
+    .unwrap()
+}

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

@@ -781,3 +781,37 @@ fn storage_is_cleared_completely_after_pull_lazy() {
     })
     .unwrap()
 }
+
+#[test]
+#[should_panic(expected = "storage entry was empty")]
+fn drop_works() {
+    ink_env::test::run_test::<ink_env::DefaultEnvironment, _>(|_| {
+        let root_key = Key::from([0x42; 32]);
+
+        // if the setup panics it should not cause the test to pass
+        let setup_result = std::panic::catch_unwind(|| {
+            let stash = create_holey_stash();
+            SpreadLayout::push_spread(&stash, &mut KeyPtr::from(root_key));
+            let _ = <StorageStash<u8> as SpreadLayout>::pull_spread(&mut KeyPtr::from(
+                root_key,
+            ));
+            // stash is dropped which should clear the cells
+        });
+        assert!(setup_result.is_ok(), "setup should not panic");
+
+        let contract_id = ink_env::test::get_current_contract_account_id::<
+            ink_env::DefaultEnvironment,
+        >()
+        .expect("Cannot get contract id");
+        let used_cells = ink_env::test::count_used_storage_cells::<
+            ink_env::DefaultEnvironment,
+        >(&contract_id)
+        .expect("used cells must be returned");
+        assert_eq!(used_cells, 0);
+
+        let _ =
+            <StorageStash<u8> as SpreadLayout>::pull_spread(&mut KeyPtr::from(root_key));
+        Ok(())
+    })
+    .unwrap()
+}