Fixed #234

dotnet · May 14, 2024 · 5ec53e6 · 5ec53e6
1 parent a39ab3a
commit 5ec53e6
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 29 deletions.
diff --git a/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachine.Pooling.cs b/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachine.Pooling.cs
@@ -136,17 +136,16 @@ void IAsyncStateMachine.MoveNext()
         }
 
         // finalize state machine
-        if (StateId == IAsyncStateMachine<TState>.FinalState)
+        if (StateId is IAsyncStateMachine<TState>.FinalState)
         {
-            if (exception is null)
-                builder.SetResult();
+            if (exception?.SourceException is { } e)
+            {
+                builder.SetException(e);
+            }
             else
-                builder.SetException(exception.SourceException);
-
-            // perform cleanup after resuming of all suspended tasks
-            guardedRegionsCounter = 0;
-            exception = null;
-            State = default;
+            {
+                builder.SetResult();
+            }
         }
     }
 
@@ -376,18 +375,16 @@ void IAsyncStateMachine.MoveNext()
         }
 
         // finalize state machine
-        if (StateId == IAsyncStateMachine<TState>.FinalState)
+        if (StateId is IAsyncStateMachine<TState>.FinalState)
         {
-            if (exception is null)
-                builder.SetResult(result);
+            if (exception?.SourceException is { } e)
+            {
+                builder.SetException(e);
+            }
             else
-                builder.SetException(exception.SourceException);
-
-            // perform cleanup after resuming of all suspended tasks
-            guardedRegionsCounter = 0;
-            exception = null;
-            result = default;
-            State = default;
+            {
+                builder.SetResult(result);
+            }
         }
     }
 

diff --git a/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachine.cs b/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachine.cs
@@ -136,12 +136,16 @@ void IAsyncStateMachine.MoveNext()
         }
 
         // finalize state machine
-        if (StateId == IAsyncStateMachine<TState>.FinalState)
+        if (StateId is IAsyncStateMachine<TState>.FinalState)
         {
-            if (exception is null)
-                builder.SetResult();
+            if (exception?.SourceException is { } e)
+            {
+                builder.SetException(e);
+            }
             else
-                builder.SetException(exception.SourceException);
+            {
+                builder.SetResult();
+            }
 
             // perform cleanup after resuming of all suspended tasks
             guardedRegionsCounter = 0;
@@ -376,12 +380,16 @@ void IAsyncStateMachine.MoveNext()
         }
 
         // finalize state machine
-        if (StateId == IAsyncStateMachine<TState>.FinalState)
+        if (StateId is IAsyncStateMachine<TState>.FinalState)
         {
-            if (exception is null)
-                builder.SetResult(result);
+            if (exception?.SourceException is { } e)
+            {
+                builder.SetException(e);
+            }
             else
-                builder.SetException(exception.SourceException);
+            {
+                builder.SetResult(result);
+            }
 
             // perform cleanup after resuming of all suspended tasks
             guardedRegionsCounter = 0;

diff --git a/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachineBuilder.cs b/src/DotNext.Metaprogramming/Runtime/CompilerServices/AsyncStateMachineBuilder.cs
@@ -532,12 +532,12 @@ internal MemberExpression Build(Type stateType)
             Type stateMachineType;
             if (returnType == typeof(void))
             {
-                stateMachineType = usePooling ? typeof(AsyncStateMachine<>) : typeof(PoolingAsyncStateMachine<>);
+                stateMachineType = usePooling ? typeof(PoolingAsyncStateMachine<>) : typeof(AsyncStateMachine<>);
                 stateMachineType = stateMachineType.MakeGenericType(stateType);
             }
             else
             {
-                stateMachineType = usePooling ? typeof(AsyncStateMachine<,>) : typeof(PoolingAsyncStateMachine<,>);
+                stateMachineType = usePooling ? typeof(PoolingAsyncStateMachine<,>) : typeof(AsyncStateMachine<,>);
                 stateMachineType = stateMachineType.MakeGenericType(stateType, returnType);
             }
 

diff --git a/src/DotNext.Tests/Metaprogramming/LambdaTests.cs b/src/DotNext.Tests/Metaprogramming/LambdaTests.cs
@@ -475,4 +475,22 @@ public static async Task ParameterClosure()
 
         Equal("hello, world", await lambda.Compile().Invoke(", world"));
     }
+
+    [Fact]
+    public static async Task RegressionIssue234()
+    {
+        var asyncMethod = new Func<Task>(DoAsync).Method;
+        var lambda = AsyncLambda<Func<ValueTask<int>>>(usePooling: true, (ctx, result) =>
+        {
+            Await(Expression.Call(asyncMethod));
+            Assign(result, 42.Const());
+        }).Compile();
+
+        Equal(42, await lambda.Invoke());
+        Equal(42, await lambda.Invoke());
+        Equal(42, await lambda.Invoke());
+        Equal(42, await lambda.Invoke());
+
+        static Task DoAsync() => Task.Delay(1);
+    }
 }