Repository: paritytech/ink
Pull request #418 (2020-05-23 19:24:54)

Add method to finalize the dynamic storage allocator
So far it has not been possible to write back the accumulated state of the global dynamic storage allocator instance. With this PR we not only eliminate some code duplication but also provide this additional API in order to provide users with a way to command finalization of the global allocator for storage.

codecov-commenter: # [Codecov](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=h1) Report
> Merging [#418](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=desc) into [master](https://codecov.io/gh/paritytech/ink/commit/718b9f5ed650de5b882091e7860a44f99f5e1926&el=desc) will **increase** coverage by `0.46%`.
> The diff coverage is `65.11%`.

[![Impacted file tree graph](https://codecov.io/gh/paritytech/ink/pull/418/graphs/tree.svg?width=650&height=150&src=pr&token=tYyVaCZb7o)](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #418      +/-   ##
==========================================
+ Coverage   86.22%   86.69%   +0.46%     
==========================================
  Files         135      135              
  Lines        5831     5854      +23     
==========================================
+ Hits         5028     5075      +47     
+ Misses        803      779      -24     
```


| [Impacted Files](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [core/src/storage2/alloc/init.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-Y29yZS9zcmMvc3RvcmFnZTIvYWxsb2MvaW5pdC5ycw==) | `65.85% <62.85%> (-34.15%)` | :arrow_down: |
| [core/src/storage2/alloc/mod.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-Y29yZS9zcmMvc3RvcmFnZTIvYWxsb2MvbW9kLnJz) | `75.00% <66.66%> (-25.00%)` | :arrow_down: |
| [core/src/storage2/alloc/tests.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-Y29yZS9zcmMvc3RvcmFnZTIvYWxsb2MvdGVzdHMucnM=) | `100.00% <100.00%> (ø)` | |
| [core/src/storage2/collections/boxed/tests.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-Y29yZS9zcmMvc3RvcmFnZTIvY29sbGVjdGlvbnMvYm94ZWQvdGVzdHMucnM=) | `100.00% <100.00%> (ø)` | |
| [core/src/hash/accumulator.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-Y29yZS9zcmMvaGFzaC9hY2N1bXVsYXRvci5ycw==) | `92.59% <0.00%> (+14.81%)` | :arrow_up: |
| [primitives/src/hash.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-cHJpbWl0aXZlcy9zcmMvaGFzaC5ycw==) | `54.16% <0.00%> (+37.50%)` | :arrow_up: |
| [primitives/src/key.rs](https://codecov.io/gh/paritytech/ink/pull/418/diff?src=pr&el=tree#diff-cHJpbWl0aXZlcy9zcmMva2V5LnJz) | `87.69% <0.00%> (+41.53%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=footer). Last update [718b9f5...0ec154c](https://codecov.io/gh/paritytech/ink/pull/418?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/core/src/storage2/alloc/init.rs b/core/src/storage2/alloc/init.rs
index 28172d0b00..e21e4fd06e 100644
--- a/core/src/storage2/alloc/init.rs
+++ b/core/src/storage2/alloc/init.rs
@@ -13,8 +13,15 @@
 // limitations under the License.
 
 use super::DynamicAllocator;
-use crate::storage2::traits::pull_spread_root;
+use crate::storage2::traits::{
+    pull_spread_root,
+    push_spread_root,
+};
 use cfg_if::cfg_if;
+use core::{
+    mem,
+    mem::ManuallyDrop,
+};
 use ink_primitives::Key;
 
 /// The default dynamic allocator key offset.
@@ -53,8 +60,10 @@ enum DynamicAllocatorState {
     /// with the assumption that a former contract deployment has already
     /// taken place in the past.
     UninitCall,
-    /// The global instance has already been initialized successfully.
+    /// The global instance has been initialized successfully and can be used.
     Initialized(DynamicAllocator),
+    /// The global instance has been finalized and can no longer be used.
+    Finalized,
 }
 
 impl From<ContractPhase> for DynamicAllocatorState {
@@ -66,6 +75,104 @@ impl From<ContractPhase> for DynamicAllocatorState {
     }
 }
 
+impl DynamicAllocatorState {
+    /// Initializes the global dynamic storage allocator instance.
+    ///
+    /// The `phase` parameter describes for which execution phase the dynamic
+    /// storage allocator needs to be initialized since this is different
+    /// in contract instantiations and calls.
+    pub fn initialize(&mut self, phase: ContractPhase) {
+        match self {
+            DynamicAllocatorState::Initialized(_)
+                // We only perform this check on Wasm compilation to avoid
+                // some overly constrained check for the off-chain testing.
+                if cfg!(all(not(feature = "std"), target_arch = "wasm32")) =>
+            {
+                panic!(
+                    "cannot initialize the dynamic storage \
+                     allocator instance twice in Wasm",
+                )
+            }
+            DynamicAllocatorState::Finalized => {
+                panic!(
+                    "cannot initialize the dynamic storage \
+                 allocator after it has been finalized",
+                )
+            }
+            state => {
+                *state = phase.into();
+            }
+        }
+    }
+
+    /// Finalizes the global instance for the dynamic storage allocator.
+    ///
+    /// The global dynamic storage allocator must not be used after this!
+    pub fn finalize(&mut self) {
+        match self {
+            DynamicAllocatorState::Initialized(allocator) => {
+                // Push all state of the global dynamic storage allocator
+                // instance back onto the contract storage.
+                push_spread_root::<DynamicAllocator>(
+                    &allocator,
+                    &Key(DYNAMIC_ALLOCATOR_KEY_OFFSET),
+                );
+                // Prevent calling `drop` on the dynamic storage allocator
+                // instance since this would clear all contract storage
+                // again.
+                let _ = ManuallyDrop::new(mem::take(allocator));
+                *self = DynamicAllocatorState::Finalized;
+            }
+            DynamicAllocatorState::Finalized => {
+                panic!(
+                    "cannot finalize the dynamic storage allocator \
+                 after it has already been finalized"
+                )
+            }
+            DynamicAllocatorState::UninitCall | DynamicAllocatorState::UninitDeploy => {
+                // Nothing to do in these states.
+            }
+        }
+    }
+
+    /// Runs the closure on the global instance for the dynamic storage allocator.
+    ///
+    /// Will automatically initialize the global allocator instance if it has not
+    /// yet been initialized.
+    ///
+    /// # Panics
+    ///
+    /// If the global dynamic storage allocator instance has already been finalized.
+    pub fn on_instance<F, R>(&mut self, f: F) -> R
+    where
+        F: FnOnce(&mut DynamicAllocator) -> R,
+    {
+        match self {
+            DynamicAllocatorState::UninitDeploy => {
+                let mut allocator = DynamicAllocator::default();
+                let result = f(&mut allocator);
+                *self = DynamicAllocatorState::Initialized(allocator);
+                result
+            }
+            DynamicAllocatorState::UninitCall => {
+                let mut allocator = pull_spread_root::<DynamicAllocator>(&Key(
+                    DYNAMIC_ALLOCATOR_KEY_OFFSET,
+                ));
+                let result = f(&mut allocator);
+                *self = DynamicAllocatorState::Initialized(allocator);
+                result
+            }
+            DynamicAllocatorState::Initialized(ref mut allocator) => f(allocator),
+            DynamicAllocatorState::Finalized => {
+                panic!(
+                    "cannot operate on the dynamic storage \
+                     allocator after it has been finalized"
+                );
+            }
+        }
+    }
+}
+
 cfg_if! {
     if #[cfg(all(not(feature = "std"), target_arch = "wasm32"))] {
         // Procedures for the Wasm compilation:
@@ -73,40 +180,36 @@ cfg_if! {
         /// The global instance for the dynamic storage allocator.
         static mut GLOBAL_INSTANCE: DynamicAllocatorState = DynamicAllocatorState::UninitDeploy;
 
-        /// Commands the (re-)initialization of the global instance for the dynamic
-        /// storage allocator.
-        pub fn initialize_for(phase: ContractPhase) {
-            let instance = unsafe { &mut GLOBAL_INSTANCE };
-            // We do not allow reinitialization for Wasm targets for performance reasons.
-            if let DynamicAllocatorState::Initialized(_) = instance {
-                panic!("cannot reinitialize dynamic storage allocator instance in Wasm");
-            }
-            *instance = phase.into();
+        /// Forwards to the `initialize` of the global dynamic storage allocator instance.
+        pub fn initialize(phase: ContractPhase) {
+            // SAFETY: Accessing the global allocator in Wasm mode is single
+            //         threaded and will not return back a reference to its
+            //         internal state. Also the `initialize` method won't
+            //         re-enter the dynamic storage in any possible way.
+            unsafe { &mut GLOBAL_INSTANCE }.initialize(phase);
+        }
+
+        /// Forwards to the `finalize` of the global dynamic storage allocator instance.
+        pub fn finalize() {
+            // SAFETY: Accessing the global allocator in Wasm mode is single
+            //         threaded and will not return back a reference to its
+            //         internal state. Also the `finalize` method won't
+            //         re-enter the dynamic storage in any possible way.
+            unsafe { &mut GLOBAL_INSTANCE }.finalize();
         }
 
-        /// Runs the given closure on the global instance for the dynamic storage allocator.
-        pub fn on_call<F, R>(f: F) -> R
+        /// Forwards to the `on_instance` of the global dynamic storage allocator instance.
+        pub fn on_instance<F, R>(f: F) -> R
         where
             F: FnOnce(&mut DynamicAllocator) -> R,
         {
-            let instance = unsafe { &mut GLOBAL_INSTANCE };
-            match instance {
-                DynamicAllocatorState::UninitDeploy => {
-                    let mut allocator = DynamicAllocator::default();
-                    let result = f(&mut allocator);
-                    *instance = DynamicAllocatorState::Initialized(allocator);
-                    result
-                }
-                DynamicAllocatorState::UninitCall => {
-                    let mut allocator = pull_spread_root::<DynamicAllocator>(&Key(DYNAMIC_ALLOCATOR_KEY_OFFSET));
-                    let result = f(&mut allocator);
-                    *instance = DynamicAllocatorState::Initialized(allocator);
-                    result
-                }
-                DynamicAllocatorState::Initialized(ref mut allocator) => {
-                    f(allocator)
-                }
-            }
+            // SAFETY: Accessing the global allocator in Wasm mode is single
+            //         threaded and will not return back a reference to its
+            //         internal state. Also this is an internal API only called
+            //         through `alloc` and `free` both of which do not return
+            //         anything that could allow to re-enter the dynamic storage
+            //         allocator instance.
+            unsafe { &mut GLOBAL_INSTANCE }.on_instance(f)
         }
 
     } else if #[cfg(feature = "std")] {
@@ -115,40 +218,31 @@ cfg_if! {
         use ::core::cell::RefCell;
         thread_local!(
             /// The global instance for the dynamic storage allocator.
-            static GLOBAL_INSTANCE: RefCell<DynamicAllocatorState> = RefCell::new(DynamicAllocatorState::UninitDeploy);
+            static GLOBAL_INSTANCE: RefCell<DynamicAllocatorState> = RefCell::new(
+                DynamicAllocatorState::UninitDeploy
+            );
         );
+        /// Forwards to the `initialize` of the global dynamic storage allocator instance.
+        pub fn initialize(phase: ContractPhase) {
+            GLOBAL_INSTANCE.with(|instance| {
+                instance.borrow_mut().initialize(phase)
+            });
+        }
 
-        /// Commands the (re-)initialization of the global instance for the dynamic
-        /// storage allocator.
-        pub fn initialize_for(phase: ContractPhase) {
+        /// Forwards to the `finalize` of the global dynamic storage allocator instance.
+        pub fn finalize() {
             GLOBAL_INSTANCE.with(|instance| {
-                instance.replace_with(|_| phase.into())
+                instance.borrow_mut().finalize()
             });
         }
 
-        /// Runs the given closure on the global instance for the dynamic storage allocator.
-        pub fn on_call<F, R>(f: F) -> R
+        /// Forwards to the `on_instance` of the global dynamic storage allocator instance.
+        pub fn on_instance<F, R>(f: F) -> R
         where
             F: FnOnce(&mut DynamicAllocator) -> R,
         {
             GLOBAL_INSTANCE.with(|instance| {
-                match &mut *instance.borrow_mut() {
-                    instance @ DynamicAllocatorState::UninitDeploy => {
-                        let mut allocator = DynamicAllocator::default();
-                        let result = f(&mut allocator);
-                        *instance = DynamicAllocatorState::Initialized(allocator);
-                        result
-                    }
-                    instance @ DynamicAllocatorState::UninitCall => {
-                        let mut allocator = pull_spread_root::<DynamicAllocator>(&Key(DYNAMIC_ALLOCATOR_KEY_OFFSET));
-                        let result = f(&mut allocator);
-                        *instance = DynamicAllocatorState::Initialized(allocator);
-                        result
-                    }
-                    DynamicAllocatorState::Initialized(instance) => {
-                        f(instance)
-                    }
-                }
+                instance.borrow_mut().on_instance(f)
             })
         }
 
diff --git a/core/src/storage2/alloc/mod.rs b/core/src/storage2/alloc/mod.rs
index 7ad3f03d04..f8cd36e662 100644
--- a/core/src/storage2/alloc/mod.rs
+++ b/core/src/storage2/alloc/mod.rs
@@ -75,21 +75,15 @@ mod init;
 #[cfg(test)]
 mod tests;
 
+use self::allocator::DynamicAllocator;
 pub use self::{
     allocation::DynamicAllocation,
-    init::{
-        initialize_for,
-        ContractPhase,
-    },
-};
-use self::{
-    allocator::DynamicAllocator,
-    init::on_call,
+    init::ContractPhase,
 };
 
 /// Returns a new dynamic storage allocation.
 pub fn alloc() -> DynamicAllocation {
-    on_call(DynamicAllocator::alloc)
+    init::on_instance(DynamicAllocator::alloc)
 }
 
 /// Frees the given dynamic storage allocation.
@@ -97,5 +91,43 @@ pub fn alloc() -> DynamicAllocation {
 /// This makes the given dynamic storage allocation available again
 /// for new dynamic storage allocations.
 pub fn free(allocation: DynamicAllocation) {
-    on_call(|allocator| allocator.free(allocation))
+    init::on_instance(|allocator| allocator.free(allocation))
+}
+
+/// Tells the global dynamic storage allocator instance how it shall initialize.
+///
+/// # Note
+///
+/// Normally users of ink! do not have to call this function directly as it is
+/// automatically being use in the correct order and way by the generated code.
+///
+/// - The `phase` parameter describes for which execution phase the dynamic
+///   storage allocator needs to be initialized since this is different
+///   in contract instantiations and calls.
+/// - This has to be issued before the first interaction with the global allocator.
+/// - The actual instantiation will happen only upon the first interaction with
+///   the global allocator, e.g. using its `alloc` or `free` calls. Until then
+///   it remains uninitialized.
+///
+/// If this function is not called before the first global allocator interaction
+/// then the default initialization scheme is for contract instantiation.
+/// However, this behavior might change and must not be relied upon.
+pub fn initialize(phase: ContractPhase) {
+    init::initialize(phase);
+}
+
+/// Finalizes the global dynamic storage allocator instance.
+///
+/// This pushes all the accumulated state from this contract execution back to
+/// the contract storage to be used in the next contract execution for the same
+/// contract instance.
+///
+/// The global dynamic storage allocator must not be used after this!
+///
+/// # Note
+///
+/// Normally users of ink! do not have to call this function directly as it is
+/// automatically being use in the correct order and way by the generated code.
+pub fn finalize() {
+    init::finalize()
 }


Test patch:
diff --git a/core/src/storage2/alloc/tests.rs b/core/src/storage2/alloc/tests.rs
index 64908d65bf..c8731463f9 100644
--- a/core/src/storage2/alloc/tests.rs
+++ b/core/src/storage2/alloc/tests.rs
@@ -15,7 +15,6 @@
 use super::{
     alloc,
     free,
-    initialize_for,
     ContractPhase,
     DynamicAllocation,
     DynamicAllocator,
@@ -25,9 +24,12 @@ use crate::{
         test,
         DefaultEnvTypes,
     },
-    storage2::traits::{
-        KeyPtr,
-        SpreadLayout,
+    storage2::{
+        alloc,
+        traits::{
+            KeyPtr,
+            SpreadLayout,
+        },
     },
 };
 use ink_primitives::Key;
@@ -36,7 +38,7 @@ fn run_default_test<F>(f: F)
 where
     F: FnOnce(),
 {
-    initialize_for(ContractPhase::Deploy);
+    alloc::initialize(ContractPhase::Deploy);
     test::run_test::<DefaultEnvTypes, _>(|_| {
         f();
         Ok(())
@@ -195,7 +197,7 @@ fn test_call_setup_works() {
         assert_eq!(allocator.alloc(), DynamicAllocation(1));
         let root_key = Key([0xFE; 32]);
         SpreadLayout::push_spread(&allocator, &mut KeyPtr::from(root_key));
-        initialize_for(ContractPhase::Call);
+        alloc::initialize(ContractPhase::Call);
         assert_eq!(alloc(), DynamicAllocation(2));
         assert_eq!(alloc(), DynamicAllocation(3));
         free(DynamicAllocation(0));
diff --git a/core/src/storage2/collections/boxed/tests.rs b/core/src/storage2/collections/boxed/tests.rs
index 04933cf62b..a331681d39 100644
--- a/core/src/storage2/collections/boxed/tests.rs
+++ b/core/src/storage2/collections/boxed/tests.rs
@@ -17,10 +17,8 @@ use crate::{
     env,
     env::test::DefaultAccounts,
     storage2::{
-        alloc::{
-            initialize_for,
-            ContractPhase,
-        },
+        alloc,
+        alloc::ContractPhase,
         traits::{
             KeyPtr,
             SpreadLayout,
@@ -50,7 +48,7 @@ where
     F: FnOnce(DefaultAccounts<env::DefaultEnvTypes>),
 {
     env::test::run_test::<env::DefaultEnvTypes, _>(|default_accounts| {
-        initialize_for(ContractPhase::Deploy);
+        alloc::initialize(ContractPhase::Deploy);
         f(default_accounts);
         Ok(())
     })