From 3d6620b3092aa09a36c57ab423f4d9db82e806cc Mon Sep 17 00:00:00 2001
From: Andy Lok <andylokandy@hotmail.com>
Date: Sat, 22 Jul 2023 15:10:53 +0800
Subject: [PATCH] fix return from block

---
 Cargo.toml                 |  2 +-
 src/lib.rs                 | 21 ++++++++++++---------
 tests/ui/ok/unreachable.rs | 13 +++++++++++++
 3 files changed, 26 insertions(+), 10 deletions(-)
 create mode 100644 tests/ui/ok/unreachable.rs

diff --git a/Cargo.toml b/Cargo.toml
index 110d3f4..1de35cd 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "logcall"
-version = "0.1.2"
+version = "0.1.3"
 edition = "2021"
 authors = ["andylokandy <andylokandy@hotmail.com>"]
 description = "An attribute macro that logs the return value from function call."
diff --git a/src/lib.rs b/src/lib.rs
index f9188cf..41f1b80 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -165,19 +165,19 @@ fn gen_block(
     args: Args,
 ) -> proc_macro2::TokenStream {
     match args {
-        Args::Simple { level } =>  {
+        Args::Simple { level } => {
             // Generate the instrumented function body.
             // If the function is an `async fn`, this will wrap it in an async block.
             if async_context {
                 let log = gen_log(&level, fn_name, "__ret_value");
                 let block = quote_spanned!(block.span()=>
                     async move {
-                        let __ret_value = #block;
+                        let __ret_value = async move { #block }.await;
                         #log;
                         __ret_value
                     }
                 );
-        
+
                 if async_keyword {
                     quote_spanned!(block.span()=>
                         #block.await
@@ -188,13 +188,16 @@ fn gen_block(
             } else {
                 let log = gen_log(&level, fn_name, "__ret_value");
                 quote_spanned!(block.span()=>
-                    let __ret_value = #block;
+                    let __ret_value = (move || #block)();
                     #log;
                     __ret_value
                 )
             }
         }
-        Args::Result { ok_level, err_level } => {
+        Args::Result {
+            ok_level,
+            err_level,
+        } => {
             let ok_arm = if let Some(ok_level) = ok_level {
                 let log_ok = gen_log(&ok_level, fn_name, "__ret_value");
                 quote_spanned!(block.span()=>
@@ -221,20 +224,20 @@ fn gen_block(
                     Err(__ret_value) => Err(__ret_value),
                 )
             };
-            
+
             // Generate the instrumented function body.
             // If the function is an `async fn`, this will wrap it in an async block.
             if async_context {
                 let block = quote_spanned!(block.span()=>
                     async move {
-                        let __ret_value = #block;
+                        let __ret_value = async move { #block }.await;
                         match __ret_value {
                             #ok_arm
                             #err_arm
                         }
                     }
                 );
-        
+
                 if async_keyword {
                     quote_spanned!(block.span()=>
                         #block.await
@@ -244,7 +247,7 @@ fn gen_block(
                 }
             } else {
                 quote_spanned!(block.span()=>
-                    let __ret_value = #block;
+                    let __ret_value = (move || #block)();
                     match __ret_value {
                         #ok_arm
                         #err_arm
diff --git a/tests/ui/ok/unreachable.rs b/tests/ui/ok/unreachable.rs
new file mode 100644
index 0000000..c311896
--- /dev/null
+++ b/tests/ui/ok/unreachable.rs
@@ -0,0 +1,13 @@
+#[logcall::logcall("info")]
+async fn f(a: u32) -> u32 {
+    if a == 1 {
+        return 1;
+    }
+
+    unreachable!()    
+}
+
+#[tokio::main]
+async fn main() {
+    f(1).await;
+}