From e0d8410b627ccfb8e751f0283404622ae3a74b14 Mon Sep 17 00:00:00 2001
From: neuecc <ils@neue.cc>
Date: Tue, 15 Sep 2020 06:35:22 +0900
Subject: [PATCH] fix IEnumerator.ToUniTask, does not work correctly when
 coroutine complete immediately

---
 .../Runtime/EnumeratorAsyncExtensions.cs      | 25 +++++++++++-
 .../Assets/Tests/CoroutineToUniTaskTest.cs    | 39 +++++++++++++++++++
 2 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/EnumeratorAsyncExtensions.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/EnumeratorAsyncExtensions.cs
index c2e4c51..7e3ebe7 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/EnumeratorAsyncExtensions.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/EnumeratorAsyncExtensions.cs
@@ -46,6 +46,8 @@ namespace Cysharp.Threading.Tasks
             IEnumerator innerEnumerator;
             CancellationToken cancellationToken;
             int initialFrame;
+            bool loopRunning;
+            bool calledGetResult;
 
             UniTaskCompletionSourceCore<object> core;
 
@@ -68,6 +70,8 @@ namespace Cysharp.Threading.Tasks
 
                 result.innerEnumerator = ConsumeEnumerator(innerEnumerator);
                 result.cancellationToken = cancellationToken;
+                result.loopRunning = true;
+                result.calledGetResult = false;
                 result.initialFrame = -1;
 
                 PlayerLoopHelper.AddAction(timing, result);
@@ -82,11 +86,15 @@ namespace Cysharp.Threading.Tasks
             {
                 try
                 {
+                    calledGetResult = true;
                     core.GetResult(token);
                 }
                 finally
                 {
-                    TryReturn();
+                    if (!loopRunning)
+                    {
+                        TryReturn();
+                    }
                 }
             }
 
@@ -107,8 +115,21 @@ namespace Cysharp.Threading.Tasks
 
             public bool MoveNext()
             {
+                if (calledGetResult)
+                {
+                    loopRunning = false;
+                    TryReturn();
+                    return false;
+                }
+
+                if (innerEnumerator == null) // invalid status, returned but loop running?
+                {
+                    return false;
+                }
+
                 if (cancellationToken.IsCancellationRequested)
                 {
+                    loopRunning = false;
                     core.TrySetCanceled(cancellationToken);
                     return false;
                 }
@@ -135,10 +156,12 @@ namespace Cysharp.Threading.Tasks
                 }
                 catch (Exception ex)
                 {
+                    loopRunning = false;
                     core.TrySetException(ex);
                     return false;
                 }
 
+                loopRunning = false;
                 core.TrySetResult(null);
                 return false;
             }
diff --git a/src/UniTask/Assets/Tests/CoroutineToUniTaskTest.cs b/src/UniTask/Assets/Tests/CoroutineToUniTaskTest.cs
index 6f48185..a27f30b 100644
--- a/src/UniTask/Assets/Tests/CoroutineToUniTaskTest.cs
+++ b/src/UniTask/Assets/Tests/CoroutineToUniTaskTest.cs
@@ -96,6 +96,45 @@ namespace Cysharp.Threading.TasksTests
         //    l[1].Item2.Should().NotBe(currentFrame);
         //}
 
+        [UnityTest]
+        public IEnumerator ImmediateRunTest() => UniTask.ToCoroutine(async () =>
+        {
+            var l = new List<int>();
+            var x1 = Immediate(l);
+            var x2 = Immediate(l);
+            var x3 = DelayOne(l);
+
+            var t1 = x1.ToUniTask();
+            CollectionAssert.AreEqual(l, new[] { 1, 2, 3 });
+            await t1;
+
+            var t2 = x2.ToUniTask();
+            CollectionAssert.AreEqual(l, new[] { 1, 2, 3, 1, 2, 3 });
+
+            var t3 = x3.ToUniTask();
+            CollectionAssert.AreEqual(l, new[] { 1, 2, 3, 1, 2, 3, 10 });
+
+            await UniTask.WhenAll(t2, t3);
+            CollectionAssert.AreEqual(l, new[] { 1, 2, 3, 1, 2, 3, 10, 20, 30 });
+
+        });
+
+        IEnumerator Immediate(List<int> l)
+        {
+            l.Add(1);
+            l.Add(2);
+            l.Add(3);
+            yield break;
+        }
+
+        IEnumerator DelayOne(List<int> l)
+        {
+            l.Add(10);
+            yield return null;
+            l.Add(20);
+            l.Add(30);
+        }
+
         [UnityTest]
         public IEnumerator WaitForSecondsTest() => UniTask.ToCoroutine(async () =>
         {