From a455de88b00d0b0f12e9734b46831d9c04b5054b Mon Sep 17 00:00:00 2001
From: RamType0 <hiromu@mt-planning.com>
Date: Thu, 17 Dec 2020 14:28:10 +0900
Subject: [PATCH] Suppress allocations around UniTask<T>.AsUniTask and
 UniTask.AsAsyncUnitUniTask

- Fix completedTask.AsUniTask or completedTask.AsAsyncUnitUniTask leaks IUniTaskSource
- UniTask.AsAsyncUnitUniTask no longer allocates when its inner IUniTaskSource is IUniTaskSource<AsyncUnit>
---
 src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.cs | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.cs b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.cs
index 26aa78a..8eb2087 100644
--- a/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.cs
+++ b/src/UniTask/Assets/Plugins/UniTask/Runtime/UniTask.cs
@@ -115,8 +115,13 @@ namespace Cysharp.Threading.Tasks
             var status = this.source.GetStatus(this.token);
             if (status.IsCompletedSuccessfully())
             {
+                this.source.GetResult(this.token);
                 return CompletedTasks.AsyncUnit;
             }
+            else if(this.source is IUniTaskSource<AsyncUnit> asyncUnitSource)
+            {
+                return new UniTask<AsyncUnit>(asyncUnitSource, this.token);
+            }
 
             return new UniTask<AsyncUnit>(new AsyncUnitSource(this.source), this.token);
         }
@@ -422,6 +427,7 @@ namespace Cysharp.Threading.Tasks
             var status = this.source.GetStatus(this.token);
             if (status.IsCompletedSuccessfully())
             {
+                this.source.GetResult(this.token);
                 return UniTask.CompletedTask;
             }