Fixed #223

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);
 }

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));
     }

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");
+        }
+    }
+}