Follow-up: SpreadLayout (#995)

* implement const_fn Key::new constructor

* add ContractRootKey trait to ink_lang crate

* use qualified path

* re-export ContractRootKey trait

* automatically implement ContractRootKey for ink! contracts

* implement #[derive(SpreadAllocate)] macro & tests

* re-export #[derive(SpreadAllocate)] macro from ink_storage::traits

* automatically implement SpreadAllocate for the ink! storage struct

* extend execute_constructor to work with fallible ink! constructors

Note that this commit does not yet allow to write fallible ink! constructors.

* add initialize_contract utility method

* re-export initialize_contract utility method

* adjust ink! codegen for recent changes with execute_constructor

* fix clippy warning

* fix hunspell dictionary

* add missing docs to codegen types

* fix UI test

* remove automatic SpreadAllocate implementation

This is due to the fact that trivial_bounds feature is not yet stable.
Otherwise this automatically generated implementation would be totally fine.
We can re-add this feature once Rust adds trivial trait bounds.

* Satisfy OCD in test

* Remove mention of `enum`s from `derive_struct()`

* Remove outdated comment about Clippy lint

* RustFmt

Co-authored-by: Hernando Castano <hernando@hcastano.com>

.config/cargo_spellcheck.dic

@@ -103,3 +103,4 @@ wildcard/S
 natively
 payability
 unpayable
+initializer

crates/lang/codegen/src/generator/dispatch.rs

@@ -511,7 +511,7 @@ impl Dispatch<'_> {
                 .is_dynamic_storage_allocator_enabled();
             quote_spanned!(constructor_span=>
                 Self::#constructor_ident(input) => {
-                    ::ink_lang::codegen::execute_constructor::<#storage_ident, _>(
+                    ::ink_lang::codegen::execute_constructor::<#storage_ident, _, _>(
                         ::ink_lang::codegen::ExecuteConstructorConfig {
                             dynamic_storage_alloc: #is_dynamic_storage_allocation_enabled
                         },

crates/lang/codegen/src/generator/storage.rs

@@ -39,8 +39,8 @@ impl GenerateCode for Storage<'_> {
                 quote! { use ::ink_lang::codegen::EmitEvent as _; }
             });
         quote_spanned!(storage_span =>
-            #access_env_impls
             #storage_struct
+            #access_env_impls
 
             const _: () = {
                 // Used to make `self.env()` and `Self::env()` available in message code.
@@ -106,6 +106,10 @@ impl Storage<'_> {
                 impl ::ink_lang::reflect::ContractName for #ident {
                     const NAME: &'static str = ::core::stringify!(#ident);
                 }
+
+                impl ::ink_lang::codegen::ContractRootKey for #ident {
+                    const ROOT_KEY: ::ink_primitives::Key = ::ink_primitives::Key::new([0x00; 32]);
+                }
             };
         )
     }

crates/lang/src/codegen/dispatch/execution.rs

@@ -18,23 +18,41 @@ use crate::reflect::{
 };
 use core::{
     any::TypeId,
+    convert::Infallible,
     mem::ManuallyDrop,
 };
 use ink_env::{
     Environment,
     ReturnFlags,
 };
-use ink_primitives::Key;
+use ink_primitives::{
+    Key,
+    KeyPtr,
+};
 use ink_storage::{
     alloc,
     alloc::ContractPhase,
     traits::{
         pull_spread_root,
         push_spread_root,
+        SpreadAllocate,
         SpreadLayout,
     },
 };
 
+/// The root key of the ink! smart contract.
+///
+/// # Note
+///
+/// - This is the key where storage allocation, pushing and pulling is rooted
+///   using the `SpreadLayout` and `SpreadAllocate` traits primarily.
+/// - This trait is automatically implemented by the ink! codegen.
+/// - The existence of this trait allows to customize the root key in future
+///   versions of ink! if needed.
+pub trait ContractRootKey {
+    const ROOT_KEY: Key;
+}
+
 /// Returns `Ok` if the caller did not transfer additional value to the callee.
 ///
 /// # Errors
@@ -70,24 +88,192 @@ pub struct ExecuteConstructorConfig {
 /// The closure is supposed to already contain all the arguments that the real
 /// constructor message requires and forwards them.
 #[inline]
-pub fn execute_constructor<S, F>(
+pub fn execute_constructor<Contract, F, R>(
     config: ExecuteConstructorConfig,
     f: F,
 ) -> Result<(), DispatchError>
 where
-    S: ink_storage::traits::SpreadLayout,
-    F: FnOnce() -> S,
+    Contract: SpreadLayout + ContractRootKey,
+    F: FnOnce() -> R,
+    <private::Seal<R> as ConstructorReturnType<Contract>>::ReturnValue: scale::Encode,
+    private::Seal<R>: ConstructorReturnType<Contract>,
 {
     if config.dynamic_storage_alloc {
         alloc::initialize(ContractPhase::Deploy);
     }
-    let storage = ManuallyDrop::new(f());
-    let root_key = Key::from([0x00; 32]);
-    push_spread_root::<S>(&storage, &root_key);
-    if config.dynamic_storage_alloc {
-        alloc::finalize();
+    let result = ManuallyDrop::new(private::Seal(f()));
+    match result.as_result() {
+        Ok(contract) => {
+            // Constructor is infallible or is fallible but succeeded.
+            //
+            // This requires us to sync back the changes of the contract storage.
+            let root_key = <Contract as ContractRootKey>::ROOT_KEY;
+            push_spread_root::<Contract>(contract, &root_key);
+            if config.dynamic_storage_alloc {
+                alloc::finalize();
+            }
+            Ok(())
+        }
+        Err(_) => {
+            // Constructor is fallible and failed.
+            //
+            // We need to revert the state of the transaction.
+            ink_env::return_value::<
+                <private::Seal<R> as ConstructorReturnType<Contract>>::ReturnValue,
+            >(
+                ReturnFlags::default().set_reverted(true),
+                result.return_value(),
+            )
+        }
+    }
+}
+
+/// Initializes the ink! contract using the given initialization routine.
+///
+/// # Note
+///
+/// - This uses `SpreadAllocate` trait in order to default initialize the
+///   ink! smart contract before calling the initialization routine.
+/// - This either returns `Contract` or `Result<Contract, E>` depending
+///   on the return type `R` of the initializer closure `F`.
+///   If `R` is `()` then `Contract` is returned and if `R` is any type of
+///   `Result<(), E>` then `Result<Contract, E>` is returned.
+///   Other return types for `F` than the ones listed above are not allowed.
+#[inline]
+pub fn initialize_contract<Contract, F, R>(
+    initializer: F,
+) -> <R as InitializerReturnType<Contract>>::Wrapped
+where
+    Contract: ContractRootKey + SpreadAllocate,
+    F: FnOnce(&mut Contract) -> R,
+    R: InitializerReturnType<Contract>,
+{
+    let mut key_ptr = KeyPtr::from(<Contract as ContractRootKey>::ROOT_KEY);
+    let mut instance = <Contract as SpreadAllocate>::allocate_spread(&mut key_ptr);
+    let result = initializer(&mut instance);
+    result.into_wrapped(instance)
+}
+
+mod private {
+    /// Seals the implementation of `ContractInitializerReturnType`.
+    pub trait Sealed {}
+    impl Sealed for () {}
+    impl<T, E> Sealed for Result<T, E> {}
+    /// A thin-wrapper type that automatically seals its inner type.
+    ///
+    /// Since it is private it can only be used from within this crate.
+    /// We need this type in order to properly seal the `ConstructorReturnType`
+    /// trait from unwanted external trait implementations.
+    #[repr(transparent)]
+    pub struct Seal<T>(pub T);
+    impl<T> Sealed for Seal<T> {}
+}
+
+/// Guards against using invalid contract initializer types.
+///
+/// # Note
+///
+/// Currently the only allowed types are `()` and `Result<(), E>`
+/// where `E` is some unspecified error type.
+/// If the contract initializer returns `Result::Err` the utility
+/// method that is used to initialize an ink! smart contract will
+/// revert the state of the contract instantiation.
+pub trait ConstructorReturnType<C>: private::Sealed {
+    /// Is `true` if `Self` is `Result<C, E>`.
+    const IS_RESULT: bool = false;
+
+    /// The error type of the constructor return type.
+    ///
+    /// # Note
+    ///
+    /// For infallible constructors this is `core::convert::Infallible`.
+    type Error;
+
+    /// The type of the return value of the constructor.
+    ///
+    /// # Note
+    ///
+    /// For infallible constructors this is `()` whereas for fallible
+    /// constructors this is the actual return value. Since we only ever
+    /// return a value in case of `Result::Err` the `Result::Ok` value
+    /// does not matter.
+    type ReturnValue;
+
+    /// Converts the return value into a `Result` instance.
+    ///
+    /// # Note
+    ///
+    /// For infallible constructor returns this always yields `Ok`.
+    fn as_result(&self) -> Result<&C, &Self::Error>;
+
+    /// Returns the actual return value of the constructor.
+    ///
+    /// # Note
+    ///
+    /// For infallible constructor returns this always yields `()`
+    /// and is basically ignored since this does not get called
+    /// if the constructor did not fail.
+    fn return_value(&self) -> &Self::ReturnValue;
+}
+
+impl<C> ConstructorReturnType<C> for private::Seal<C> {
+    type Error = Infallible;
+    type ReturnValue = ();
+
+    #[inline(always)]
+    fn as_result(&self) -> Result<&C, &Self::Error> {
+        Ok(&self.0)
+    }
+
+    #[inline(always)]
+    fn return_value(&self) -> &Self::ReturnValue {
+        &()
+    }
+}
+
+impl<C, E> ConstructorReturnType<C> for private::Seal<Result<C, E>> {
+    const IS_RESULT: bool = true;
+    type Error = E;
+    type ReturnValue = Result<C, E>;
+
+    #[inline(always)]
+    fn as_result(&self) -> Result<&C, &Self::Error> {
+        self.0.as_ref()
+    }
+
+    #[inline(always)]
+    fn return_value(&self) -> &Self::ReturnValue {
+        &self.0
+    }
+}
+
+/// Trait used to convert return types of contract initializer routines.
+///
+/// Only `()` and `Result<(), E>` are allowed contract initializer return types.
+/// For `WrapReturnType<C>` where `C` is the contract type the trait converts
+/// `()` into `C` and `Result<(), E>` into `Result<C, E>`.
+pub trait InitializerReturnType<C>: private::Sealed {
+    type Wrapped;
+
+    /// Performs the type conversion of the initialization routine return type.
+    fn into_wrapped(self, wrapped: C) -> Self::Wrapped;
+}
+
+impl<C> InitializerReturnType<C> for () {
+    type Wrapped = C;
+
+    #[inline]
+    fn into_wrapped(self, wrapped: C) -> C {
+        wrapped
+    }
+}
+impl<C, E> InitializerReturnType<C> for Result<(), E> {
+    type Wrapped = Result<C, E>;
+
+    #[inline]
+    fn into_wrapped(self, wrapped: C) -> Self::Wrapped {
+        self.map(|_| wrapped)
     }
-    Ok(())
 }
 
 /// Configuration for execution of ink! messages.

crates/lang/src/codegen/dispatch/mod.rs

@@ -21,7 +21,9 @@ pub use self::{
         deny_payment,
         execute_constructor,
         finalize_message,
+        initialize_contract,
         initiate_message,
+        ContractRootKey,
         ExecuteConstructorConfig,
         ExecuteMessageConfig,
     },

crates/lang/src/codegen/mod.rs

@@ -26,8 +26,10 @@ pub use self::{
         deny_payment,
         execute_constructor,
         finalize_message,
+        initialize_contract,
         initiate_message,
         ContractCallBuilder,
+        ContractRootKey,
         DispatchInput,
         DispatchOutput,
         ExecuteConstructorConfig,

crates/lang/tests/ui/contract/fail/message-returns-non-codec.stderr

@@ -6,9 +6,9 @@ error[E0277]: the trait bound `NonCodecType: WrapperTypeEncode` is not satisfied
    |
    = note: required because of the requirements on the impl of `Encode` for `NonCodecType`
 note: required by a bound in `DispatchOutput`
-  --> src/codegen/dispatch/type_check.rs:69:8
+  --> src/codegen/dispatch/type_check.rs
    |
-69 |     T: scale::Encode + 'static;
+   |     T: scale::Encode + 'static;
    |        ^^^^^^^^^^^^^ required by this bound in `DispatchOutput`
 
 error[E0277]: the trait bound `NonCodecType: WrapperTypeEncode` is not satisfied
@@ -21,9 +21,9 @@ error[E0277]: the trait bound `NonCodecType: WrapperTypeEncode` is not satisfied
     |
     = note: required because of the requirements on the impl of `Encode` for `NonCodecType`
 note: required by a bound in `finalize_message`
-   --> src/codegen/dispatch/execution.rs:169:8
+   --> src/codegen/dispatch/execution.rs
     |
-169 |     R: scale::Encode + 'static,
+    |     R: scale::Encode + 'static,
     |        ^^^^^^^^^^^^^ required by this bound in `finalize_message`
 
 error[E0599]: the method `fire` exists for struct `ink_env::call::CallBuilder<DefaultEnvironment, Set<ink_env::AccountId>, Unset<u64>, Unset<u128>, Set<ExecutionInput<ArgumentList<ArgumentListEnd, ArgumentListEnd>>>, Set<ReturnType<NonCodecType>>>`, but its trait bounds were not satisfied

crates/primitives/src/key.rs

@@ -46,10 +46,23 @@ use scale_info::{
 #[repr(transparent)]
 pub struct Key([u8; 32]);
 
+impl Key {
+    /// Creates a new key instance from the given bytes.
+    ///
+    /// # Note
+    ///
+    /// This constructor only exists since it is not yet possible to define
+    /// the `From` trait implementation as const.
+    #[inline]
+    pub const fn new(bytes: [u8; 32]) -> Self {
+        Self(bytes)
+    }
+}
+
 impl From<[u8; 32]> for Key {
     #[inline]
     fn from(bytes: [u8; 32]) -> Self {
-        Self(bytes)
+        Self::new(bytes)
     }
 }