Fix ink! message result return reverts

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

@@ -672,18 +672,25 @@ impl Dispatch<'_> {
                 .is_dynamic_storage_allocator_enabled();
             quote_spanned!(message_span=>
                 Self::#message_ident(input) => {
-                    ::ink_lang::codegen::execute_message::<
-                        #storage_ident,
-                        #message_output,
-                        _
-                    >(
-                        ::ink_lang::codegen::ExecuteMessageConfig {
-                            payable: #accepts_payment,
-                            mutates: #mutates_storage,
-                            dynamic_storage_alloc: #is_dynamic_storage_allocation_enabled,
-                        },
-                        move |storage: &mut #storage_ident| { #message_callable(storage, input) }
-                    )
+                    let config = ::ink_lang::codegen::ExecuteMessageConfig {
+                        payable: #accepts_payment,
+                        mutates: #mutates_storage,
+                        dynamic_storage_alloc: #is_dynamic_storage_allocation_enabled,
+                    };
+                    let mut contract: ::core::mem::ManuallyDrop<#storage_ident> =
+                        ::core::mem::ManuallyDrop::new(
+                            ::ink_lang::codegen::initiate_message::<#storage_ident>(config)?
+                        );
+                    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(())
                 }
             )
         });

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

@@ -115,48 +115,96 @@ pub struct ExecuteMessageConfig {
     pub dynamic_storage_alloc: bool,
 }
 
-/// Executes the given `&mut self` message closure.
+/// Initiates an ink! message call with the given configuration.
+///
+/// Returns the contract state pulled from the root storage region upon success.
 ///
 /// # Note
 ///
-/// The closure is supposed to already contain all the arguments that the real
-/// message requires and forwards them.
-#[inline]
-pub fn execute_message<Storage, Output, F>(
+/// 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(always)]
+pub fn initiate_message<Contract>(
     config: ExecuteMessageConfig,
-    f: F,
-) -> Result<(), DispatchError>
+) -> Result<Contract, DispatchError>
 where
-    Storage: SpreadLayout + ContractEnv,
-    Output: scale::Encode + 'static,
-    F: FnOnce(&mut Storage) -> Output,
+    Contract: SpreadLayout + ContractEnv,
 {
     if !config.payable {
-        deny_payment::<<Storage as ContractEnv>::Env>()?;
+        deny_payment::<<Contract as ContractEnv>::Env>()?;
     }
     if config.dynamic_storage_alloc {
         alloc::initialize(ContractPhase::Call);
     }
     let root_key = Key::from([0x00; 32]);
-    let mut storage = ManuallyDrop::new(pull_spread_root::<Storage>(&root_key));
-    let result = f(&mut storage);
+    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 {
-        push_spread_root::<Storage>(&storage, &root_key);
+        let root_key = Key::from([0x00; 32]);
+        push_spread_root::<Contract>(contract, &root_key);
     }
     if config.dynamic_storage_alloc {
         alloc::finalize();
     }
-    if TypeId::of::<Output>() != TypeId::of::<()>() {
-        // We include a check for `is_result_type!(Output)` despite the fact that this
-        // is indirectly covered by `is_result_err!(&result)` because the Rust compiler
-        // will have more opportunities to optimize the whole conditional away. This is
-        // due to the fact that `is_result_type!` relies on constant information whereas
-        // is_result_err!` requires `&self`.
-        let revert_state = is_result_type!(Output) && is_result_err!(&result);
-        ink_env::return_value::<Output>(
-            ReturnFlags::default().set_reverted(revert_state),
-            &result,
-        )
+    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,7 +20,8 @@ pub use self::{
     execution::{
         deny_payment,
         execute_constructor,
-        execute_message,
+        finalize_message,
+        initiate_message,
         ExecuteConstructorConfig,
         ExecuteMessageConfig,
     },

crates/lang/src/codegen/mod.rs

@@ -25,7 +25,8 @@ pub use self::{
     dispatch::{
         deny_payment,
         execute_constructor,
-        execute_message,
+        finalize_message,
+        initiate_message,
         ContractCallBuilder,
         DispatchInput,
         DispatchOutput,

crates/lang/src/lib.rs

@@ -16,7 +16,7 @@
 
 #[macro_use]
 #[doc(hidden)]
-mod result_info;
+pub mod result_info;
 
 #[cfg_attr(not(feature = "show-codegen-docs"), doc(hidden))]
 pub mod codegen;

crates/lang/src/result_info.rs

@@ -30,6 +30,8 @@ pub trait IsResultTypeFallback {
 impl<T> IsResultTypeFallback for IsResultType<T> {}
 
 /// Returns `true` if the given type is a `Result` type.
+#[macro_export]
+#[doc(hidden)]
 macro_rules! is_result_type {
     ( $T:ty $(,)? ) => {{
         #[allow(unused_imports)]
@@ -65,6 +67,8 @@ impl<T> IsResultErrFallback for IsResultErr<'_, T> {}
 /// # Note
 ///
 /// This given expression is not required to be of type `Result`.
+#[macro_export]
+#[doc(hidden)]
 macro_rules! is_result_err {
     ( $e:expr $(,)? ) => {{
         #[allow(unused_imports)]

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
@@ -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 `execute_message`
-   --> src/codegen/dispatch/execution.rs:131:13
+note: required by a bound in `finalize_message`
+   --> src/codegen/dispatch/execution.rs
     |
-131 |     Output: scale::Encode + 'static,
-    |             ^^^^^^^^^^^^^ required by this bound in `execute_message`
+    |     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
    --> tests/ui/contract/fail/message-returns-non-codec.rs:18:9