Fix lazy bug (#528)

* [storage] make Lazy::laze constructor crate private

* [storage] fix Lazy SpreadLayout impl behaving incorrectly when not used

* [storage] remove redundant state replacement in StorageEntry::push_packed_root

* [storage] add regression test for the fixed bug

* [storage] clean up code between LazyCell and StorageEntry

* [storage] apply rustfmt

crates/storage/src/lazy/entry.rs

@@ -107,20 +107,17 @@ where
 
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
         let root_key = ExtKeyPtr::next_for::<Self>(ptr);
-        Self::new(pull_spread_root_opt::<T>(&root_key), EntryState::Preserved)
+        Self::pull_spread_root(root_key)
     }
 
     fn push_spread(&self, ptr: &mut KeyPtr) {
-        let old_state = self.replace_state(EntryState::Preserved);
-        if old_state.is_mutated() {
-            let root_key = ExtKeyPtr::next_for::<Self>(ptr);
-            push_spread_root_opt::<T>(self.value().into(), &root_key);
-        }
+        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
+        self.push_spread_root(root_key)
     }
 
     fn clear_spread(&self, ptr: &mut KeyPtr) {
         let root_key = ExtKeyPtr::next_for::<Self>(ptr);
-        clear_spread_root_opt::<T, _>(&root_key, || self.value().into());
+        self.clear_spread_root(root_key)
     }
 }
 
@@ -181,6 +178,44 @@ where
     }
 }
 
+impl<T> StorageEntry<T>
+where
+    T: SpreadLayout,
+{
+    /// Pulls the entity from the underlying associated storage as spreaded representation.
+    ///
+    /// # Note
+    ///
+    /// Mainly used by lazy storage abstractions that only allow operating on
+    /// packed storage entities such as [`LazyCell`].
+    pub fn pull_spread_root(root_key: &Key) -> Self {
+        Self::new(pull_spread_root_opt::<T>(&root_key), EntryState::Preserved)
+    }
+
+    /// Pushes the underlying associated storage as spreaded representation.
+    ///
+    /// # Note
+    ///
+    /// Mainly used by lazy storage abstractions that only allow operating on
+    /// packed storage entities such as [`LazyCell`].
+    pub fn push_spread_root(&self, root_key: &Key) {
+        let old_state = self.replace_state(EntryState::Preserved);
+        if old_state.is_mutated() {
+            push_spread_root_opt::<T>(self.value().into(), &root_key);
+        }
+    }
+
+    /// Clears the underlying associated storage as spreaded representation.
+    ///
+    /// # Note
+    ///
+    /// Mainly used by lazy storage abstractions that only allow operating on
+    /// packed storage entities such as [`LazyCell`].
+    pub fn clear_spread_root(&self, root_key: &Key) {
+        clear_spread_root_opt::<T, _>(&root_key, || self.value().into());
+    }
+}
+
 impl<T> StorageEntry<T>
 where
     T: PackedLayout,
@@ -204,7 +239,6 @@ where
     pub fn push_packed_root(&self, root_key: &Key) {
         let old_state = self.replace_state(EntryState::Preserved);
         if old_state.is_mutated() {
-            self.replace_state(EntryState::Preserved);
             push_packed_root_opt::<T>(self.value().into(), &root_key);
         }
     }

crates/storage/src/lazy/lazy_cell.rs

@@ -138,18 +138,21 @@ where
     const FOOTPRINT: u64 = <T as SpreadLayout>::FOOTPRINT;
 
     fn pull_spread(ptr: &mut KeyPtr) -> Self {
-        Self::lazy(*KeyPtr::next_for::<T>(ptr))
+        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
+        Self::lazy(*root_key)
     }
 
     fn push_spread(&self, ptr: &mut KeyPtr) {
+        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
         if let Some(entry) = self.entry() {
-            SpreadLayout::push_spread(entry, ptr)
+            entry.push_spread_root(root_key)
         }
     }
 
     fn clear_spread(&self, ptr: &mut KeyPtr) {
+        let root_key = ExtKeyPtr::next_for::<Self>(ptr);
         if let Some(entry) = self.entry() {
-            SpreadLayout::clear_spread(entry, ptr)
+            entry.clear_spread_root(root_key)
         }
     }
 }
@@ -501,4 +504,45 @@ mod tests {
             Ok(())
         })
     }
+
+    #[test]
+    fn regression_test_for_issue_528() -> ink_env::Result<()> {
+        run_test::<ink_env::DefaultEnvironment, _>(|_| {
+            let root_key = Key::from([0x00; 32]);
+            {
+                // Step 1: Push a valid pair onto the contract storage.
+                let pair = (LazyCell::new(Some(1i32)), 2i32);
+                SpreadLayout::push_spread(&pair, &mut KeyPtr::from(root_key));
+            }
+            {
+                // Step 2: Pull the pair from the step before.
+                //
+                // 1. Change the second `i32` value of the pair.
+                // 2. Push the pair again to contract storage.
+                //
+                // We prevent the intermediate instance from clearing the storage preemtively by wrapping
+                // it inside `ManuallyDrop`. The third step will clean up the same storage region afterwards.
+                //
+                // We explicitely do not touch or assert the value of `pulled_pair.0` in order to trigger
+                // the bug.
+                let pulled_pair: (LazyCell<i32>, i32) =
+                    SpreadLayout::pull_spread(&mut KeyPtr::from(root_key));
+                let mut pulled_pair = core::mem::ManuallyDrop::new(pulled_pair);
+                assert_eq!(pulled_pair.1, 2i32);
+                pulled_pair.1 = 3i32;
+                SpreadLayout::push_spread(&*pulled_pair, &mut KeyPtr::from(root_key));
+            }
+            {
+                // Step 3: Pull the pair again from the storage.
+                //
+                // If the bug with `Lazy` that has been fixed in PR #528 has been fixed we should be
+                // able to inspect the correct values for both pair entries which is: `(Some(1), 3)`
+                let pulled_pair: (LazyCell<i32>, i32) =
+                    SpreadLayout::pull_spread(&mut KeyPtr::from(root_key));
+                assert_eq!(pulled_pair.0.get(), Some(&1i32));
+                assert_eq!(pulled_pair.1, 3i32);
+            }
+            Ok(())
+        })
+    }
 }

crates/storage/src/lazy/mod.rs

@@ -119,7 +119,7 @@ where
 
     /// Creates a true lazy storage value for the given key.
     #[must_use]
-    pub fn lazy(key: Key) -> Self {
+    pub(crate) fn lazy(key: Key) -> Self {
         Self {
             cell: LazyCell::lazy(key),
         }