Extracted push and pull of the contract in the dispatching logic

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

@@ -722,9 +722,8 @@ impl Dispatch<'_> {
                     }>>::IDS[#index]
                 }>>::Output
             );
-            let accepts_payment = quote_spanned!(message_span=>
-                false ||
-                <#storage_ident as ::ink_lang::reflect::DispatchableMessageInfo<{
+            let deny_payment = quote_spanned!(message_span=>
+                !<#storage_ident as ::ink_lang::reflect::DispatchableMessageInfo<{
                     <#storage_ident as ::ink_lang::reflect::ContractDispatchableMessages<{
                         <#storage_ident as ::ink_lang::reflect::ContractAmountDispatchables>::MESSAGES
                     }>>::IDS[#index]
@@ -740,24 +739,34 @@ impl Dispatch<'_> {
 
             quote_spanned!(message_span=>
                 Self::#message_ident(input) => {
-                    let config = ::ink_lang::codegen::ExecuteMessageConfig {
-                        payable: #accepts_payment,
-                        mutates: #mutates_storage,
-                    };
-                    let mut contract: ::core::mem::ManuallyDrop<#storage_ident> =
-                        ::core::mem::ManuallyDrop::new(
-                            ::ink_lang::codegen::initiate_message::<#storage_ident>(config)?
-                        );
+                    use ::core::default::Default;
+
+                    if #deny_payment {
+                        ::ink_lang::codegen::deny_payment::<
+                            <#storage_ident as ::ink_lang::reflect::ContractEnv>::Env>()?;
+                    }
+
                     let result: #message_output = #message_callable(&mut contract, input);
                     let failure = ::ink_lang::is_result_type!(#message_output)
                         && ::ink_lang::is_result_err!(result);
-                    ::ink_lang::codegen::finalize_message::<#storage_ident, #message_output>(
-                        !failure,
-                        &contract,
-                        config,
-                        &result,
-                    )?;
-                    ::core::result::Result::Ok(())
+
+                    if failure {
+                        // We return early here since there is no need to push back the
+                        // intermediate results of the contract - the transaction is going to be
+                        // reverted anyways.
+                        ::ink_env::return_value::<#message_output>(
+                            ::ink_env::ReturnFlags::default().set_reverted(true), &result
+                        )
+                    }
+
+                    push_contract(contract, #mutates_storage);
+
+                    if ::core::any::TypeId::of::<#message_output>() != ::core::any::TypeId::of::<()>() {
+                        // In case the return type is `()` we do not return a value.
+                        ::ink_env::return_value::<#message_output>(
+                            ::ink_env::ReturnFlags::default(), &result
+                        )
+                    }
                 }
             )
         });
@@ -794,12 +803,31 @@ impl Dispatch<'_> {
                     }
                 }
 
+                static ROOT_KEY: ::ink_primitives::Key = ::ink_primitives::Key::new([0x00; 32]);
+
+                fn push_contract(contract: ::core::mem::ManuallyDrop<#storage_ident>, mutates: bool) {
+                    if mutates {
+                        ::ink_storage::traits::push_spread_root::<#storage_ident>(
+                            &contract, &ROOT_KEY
+                        );
+                    }
+                }
+
                 impl ::ink_lang::reflect::ExecuteDispatchable for __ink_MessageDecoder {
                     #[allow(clippy::nonminimal_bool)]
-                    fn execute_dispatchable(self) -> ::core::result::Result<(), ::ink_lang::reflect::DispatchError> {
+                    fn execute_dispatchable(
+                        self
+                    ) -> ::core::result::Result<(), ::ink_lang::reflect::DispatchError> {
+                        let mut contract: ::core::mem::ManuallyDrop<#storage_ident> =
+                            ::core::mem::ManuallyDrop::new(
+                                ::ink_storage::traits::pull_spread_root::<#storage_ident>(&ROOT_KEY)
+                            );
+
                         match self {
                             #( #message_execute ),*
-                        }
+                        };
+
+                        ::core::result::Result::Ok(())
                     }
                 }
 

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

@@ -17,7 +17,6 @@ use crate::reflect::{
     DispatchError,
 };
 use core::{
-    any::TypeId,
     convert::Infallible,
     mem::ManuallyDrop,
 };
@@ -30,7 +29,6 @@ use ink_primitives::{
     KeyPtr,
 };
 use ink_storage::traits::{
-    pull_spread_root,
     push_spread_root,
     SpreadAllocate,
     SpreadLayout,
@@ -264,104 +262,3 @@ impl<C, E> InitializerReturnType<C> for Result<(), E> {
         self.map(|_| wrapped)
     }
 }
-
-/// Configuration for execution of ink! messages.
-#[derive(Debug, Copy, Clone)]
-pub struct ExecuteMessageConfig {
-    /// Yields `true` if the ink! message accepts payment.
-    pub payable: bool,
-    /// Yields `true` if the ink! message might mutate contract storage.
-    ///
-    /// # Note
-    ///
-    /// This is usually true for `&mut self` ink! messages.
-    pub mutates: bool,
-}
-
-/// Initiates an ink! message call with the given configuration.
-///
-/// Returns the contract state pulled from the root storage region upon success.
-///
-/// # Note
-///
-/// This work around that splits executing an ink! message into initiate
-/// and finalize phases was needed due to the fact that `is_result_type`
-/// and `is_result_err` macros do not work in generic contexts.
-#[inline]
-pub fn initiate_message<Contract>(
-    config: ExecuteMessageConfig,
-) -> Result<Contract, DispatchError>
-where
-    Contract: SpreadLayout + ContractEnv,
-{
-    if !config.payable {
-        deny_payment::<<Contract as ContractEnv>::Env>()?;
-    }
-    let root_key = Key::from([0x00; 32]);
-    let contract = pull_spread_root::<Contract>(&root_key);
-    Ok(contract)
-}
-
-/// Finalizes an ink! message call with the given configuration.
-///
-/// This dispatches into fallible and infallible message finalization
-/// depending on the given `success` state.
-///
-/// - If the message call was successful the return value is simply returned
-///   and cached storage is pushed back to the contract storage.
-/// - If the message call failed the return value result is returned instead
-///   and the transaction is signalled to be reverted.
-///
-/// # Note
-///
-/// This work around that splits executing an ink! message into initiate
-/// and finalize phases was needed due to the fact that `is_result_type`
-/// and `is_result_err` macros do not work in generic contexts.
-#[inline]
-pub fn finalize_message<Contract, R>(
-    success: bool,
-    contract: &Contract,
-    config: ExecuteMessageConfig,
-    result: &R,
-) -> Result<(), DispatchError>
-where
-    Contract: SpreadLayout,
-    R: scale::Encode + 'static,
-{
-    if success {
-        finalize_infallible_message(contract, config, result)
-    } else {
-        finalize_fallible_message(result)
-    }
-}
-
-#[inline]
-fn finalize_infallible_message<Contract, R>(
-    contract: &Contract,
-    config: ExecuteMessageConfig,
-    result: &R,
-) -> Result<(), DispatchError>
-where
-    Contract: SpreadLayout,
-    R: scale::Encode + 'static,
-{
-    if config.mutates {
-        let root_key = Key::from([0x00; 32]);
-        push_spread_root::<Contract>(contract, &root_key);
-    }
-    if TypeId::of::<R>() != TypeId::of::<()>() {
-        // In case the return type is `()` we do not return a value.
-        ink_env::return_value::<R>(ReturnFlags::default(), result)
-    }
-    Ok(())
-}
-
-#[inline]
-fn finalize_fallible_message<R>(result: &R) -> !
-where
-    R: scale::Encode + 'static,
-{
-    // There is no need to push back the intermediate results of the
-    // contract since the transaction is going to be reverted.
-    ink_env::return_value::<R>(ReturnFlags::default().set_reverted(true), result)
-}

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

@@ -20,12 +20,9 @@ pub use self::{
     execution::{
         deny_payment,
         execute_constructor,
-        finalize_message,
         initialize_contract,
-        initiate_message,
         ContractRootKey,
         ExecuteConstructorConfig,
-        ExecuteMessageConfig,
     },
     info::ContractCallBuilder,
     type_check::{

crates/lang/src/codegen/mod.rs

@@ -25,15 +25,12 @@ pub use self::{
     dispatch::{
         deny_payment,
         execute_constructor,
-        finalize_message,
         initialize_contract,
-        initiate_message,
         ContractCallBuilder,
         ContractRootKey,
         DispatchInput,
         DispatchOutput,
         ExecuteConstructorConfig,
-        ExecuteMessageConfig,
     },
     env::{
         Env,

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

@@ -20,11 +20,11 @@ error[E0277]: the trait bound `NonCodecType: WrapperTypeEncode` is not satisfied
     | |_________^ the trait `WrapperTypeEncode` is not implemented for `NonCodecType`
     |
     = 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
+note: required by a bound in `return_value`
+   --> $WORKSPACE/crates/env/src/api.rs
     |
-    |     R: scale::Encode + 'static,
-    |        ^^^^^^^^^^^^^ required by this bound in `finalize_message`
+    |     R: scale::Encode,
+    |        ^^^^^^^^^^^^^ required by this bound in `return_value`
 
 error[E0599]: the method `fire` exists for struct `ink_env::call::CallBuilder<DefaultEnvironment, Set<Call<DefaultEnvironment>>, Set<ExecutionInput<ArgumentList<ArgumentListEnd, ArgumentListEnd>>>, Set<ReturnType<NonCodecType>>>`, but its trait bounds were not satisfied
    --> tests/ui/contract/fail/message-returns-non-codec.rs:18:9