From 1f19990a2260ded469e07f92ae660aac4d2f1b2f Mon Sep 17 00:00:00 2001
From: sakno <roman.sakno@gmail.com>
Date: Sat, 17 Feb 2024 14:03:56 +0200
Subject: [PATCH] Fixed #223

---
 .../Linq/Expressions/AsyncResultExpression.cs |  3 ++
 .../AsyncStateMachineBuilder.cs               | 23 +++++++++++
 .../Metaprogramming/RegressionIssue223.cs     | 41 +++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 src/DotNext.Tests/Metaprogramming/RegressionIssue223.cs

diff --git a/src/DotNext.Metaprogramming/Linq/Expressions/AsyncResultExpression.cs b/src/DotNext.Metaprogramming/Linq/Expressions/AsyncResultExpression.cs
index 730bc4471..dbbe3eb99 100644
--- a/src/DotNext.Metaprogramming/Linq/Expressions/AsyncResultExpression.cs
+++ b/src/DotNext.Metaprogramming/Linq/Expressions/AsyncResultExpression.cs
@@ -103,4 +103,7 @@ protected override AsyncResultExpression VisitChildren(ExpressionVisitor visitor
         var expression = visitor.Visit(AsyncResult);
         return ReferenceEquals(expression, AsyncResult) ? this : new(expression, taskType);
     }
+
+    internal AsyncResultExpression Update(Expression asyncResult)
+        => new(asyncResult, taskType);
 }
\ No newline at end of file
diff --git a/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachineBuilder.cs b/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachineBuilder.cs
index b7aa4ff01..83dd8e1d8 100644
--- a/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachineBuilder.cs
+++ b/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachineBuilder.cs
@@ -60,6 +60,12 @@ internal AsyncStateMachineBuilder(Type taskType, IReadOnlyList<ParameterExpressi
         stateSwitchTable = new StateTransitionTable();
     }
 
+    internal ParameterExpression? ResultVariable
+    {
+        get;
+        private set;
+    }
+
     private static void MarkAsParameter(ParameterExpression parameter, int position)
         => parameter.GetUserData().Set(ParameterPositionSlot, position);
 
@@ -239,8 +245,18 @@ private Expression VisitAsyncResult(AsyncResultExpression expr)
         var prologue = context.CurrentStatement.PrologueCodeInserter();
         expr = (AsyncResultExpression)base.VisitExtension(expr);
 
+        var containsAwait = ExpressionAttributes.Get(expr.AsyncResult) is { ContainsAwait: true };
+
+        if (containsAwait && Task.HasResult)
+        {
+            ResultVariable ??= Expression.Parameter(Task.ResultType);
+            prologue(Expression.Assign(ResultVariable, expr.AsyncResult));
+            expr = expr.Update(ResultVariable);
+        }
+
         foreach (var finalization in context.CreateJumpPrologue(AsyncMethodEnd.Goto(), this))
             prologue(finalization);
+
         return expr;
     }
 
@@ -604,6 +620,13 @@ internal Expression<TDelegate> Build(Expression body, bool tailCall, bool usePoo
         // replace all special expressions
         body = Visit(body);
 
+        if (methodBuilder.ResultVariable is { } resultVar)
+        {
+            body = body is BlockExpression block
+                ? block.Update(block.Variables.Append(resultVar), block.Expressions)
+                : Expression.Block(body.Type, [resultVar], body);
+        }
+
         // now we have state machine method, wrap it into lambda
         return Build(BuildStateMachine(body, stateMachine, tailCall));
     }
diff --git a/src/DotNext.Tests/Metaprogramming/RegressionIssue223.cs b/src/DotNext.Tests/Metaprogramming/RegressionIssue223.cs
new file mode 100644
index 000000000..a2690d99e
--- /dev/null
+++ b/src/DotNext.Tests/Metaprogramming/RegressionIssue223.cs
@@ -0,0 +1,41 @@
+using System.Linq.Expressions;
+using System.Reflection;
+using DotNext.Linq.Expressions;
+
+namespace DotNext.Metaprogramming;
+
+using static Linq.Expressions.ExpressionBuilder;
+using static Metaprogramming.CodeGenerator;
+
+public sealed class RegressionIssue223 : Test
+{
+    [Fact]
+    public static async Task ThrowOnReturn()
+    {
+        var lambda = AsyncLambda<Func<Task<int>>>(_ =>
+        {
+            Try(() =>
+                {
+                    var methodInfo = new Func<Task<int>>(Throw).Method;
+                    var methodResult = Expression.Call(null, methodInfo);
+
+                    Return(methodResult.Await());
+                })
+                .Catch(typeof(Exception), _ =>
+                {
+                    CallStatic(typeof(Console), nameof(Console.WriteLine), Expression.Constant("Exception caught"));
+                })
+                .End();
+        });
+
+        var action = lambda.Compile();
+
+        Equal(0, await action());
+
+        static async Task<int> Throw()
+        {
+            await Task.Yield();
+            throw new InvalidOperationException("Exception was not caught");
+        }
+    }
+}
\ No newline at end of file