From 00d5c1acef96307bc58b721a274c1c37dd36ce7f Mon Sep 17 00:00:00 2001
From: Nicolas Gnyra <nicolas@gnyra.com>
Date: Wed, 14 Aug 2024 21:27:20 -0400
Subject: [PATCH] Remove animator instantiate workaround

---
 Source/CustomAvatar/Avatar/AvatarLoader.cs    | 10 +--
 .../AsyncInstantiateOperationAwaiter.cs       | 71 +++++++++++++++++++
 .../Utilities/AsyncOperationExtensions.cs     |  2 +
 3 files changed, 78 insertions(+), 5 deletions(-)
 create mode 100644 Source/CustomAvatar/Utilities/AsyncInstantiateOperationAwaiter.cs

diff --git a/Source/CustomAvatar/Avatar/AvatarLoader.cs b/Source/CustomAvatar/Avatar/AvatarLoader.cs
index 686ec149..b0685ecd 100644
--- a/Source/CustomAvatar/Avatar/AvatarLoader.cs
+++ b/Source/CustomAvatar/Avatar/AvatarLoader.cs
@@ -139,14 +139,14 @@ private async Task<AvatarPrefab> LoadAssetBundle(string fullPath, IProgress<floa
                     throw new AvatarLoadException("Could not load asset from asset bundle");
                 }
 
-                GameObject instance = UnityEngine.Object.Instantiate(prefabObject);
+                avatarPrefab = _container.InstantiateComponent<AvatarPrefab>(prefabObject, [fullPath]);
 
-                avatarPrefab = _container.InstantiateComponent<AvatarPrefab>(instance, new object[] { fullPath });
+                prefabObject.hideFlags |= HideFlags.DontUnloadUnusedAsset;
+                prefabObject.name = $"AvatarPrefab({avatarPrefab.descriptor.name})";
 
-                instance.name = $"AvatarPrefab({avatarPrefab.descriptor.name})";
-                instance.SetActive(false);
+                await ShaderRepair.FixShadersOnGameObjectAsync(prefabObject);
 
-                await ShaderRepair.FixShadersOnGameObjectAsync(instance);
+                prefabObject.hideFlags &= ~HideFlags.DontUnloadUnusedAsset;
 
                 return avatarPrefab;
             }
diff --git a/Source/CustomAvatar/Utilities/AsyncInstantiateOperationAwaiter.cs b/Source/CustomAvatar/Utilities/AsyncInstantiateOperationAwaiter.cs
new file mode 100644
index 00000000..5fac1825
--- /dev/null
+++ b/Source/CustomAvatar/Utilities/AsyncInstantiateOperationAwaiter.cs
@@ -0,0 +1,71 @@
+//  Beat Saber Custom Avatars - Custom player models for body presence in Beat Saber.
+//  Copyright © 2018-2024  Nicolas Gnyra and Beat Saber Custom Avatars Contributors
+//
+//  This library is free software: you can redistribute it and/or
+//  modify it under the terms of the GNU Lesser General Public
+//  License as published by the Free Software Foundation, either
+//  version 3 of the License, or (at your option) any later version.
+//
+//  This program is distributed in the hope that it will be useful,
+//  but WITHOUT ANY WARRANTY; without even the implied warranty of
+//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+//  GNU Lesser General Public License for more details.
+//
+//  You should have received a copy of the GNU Lesser General Public License
+//  along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+using System;
+using System.Runtime.CompilerServices;
+using UnityEngine;
+
+#pragma warning disable IDE1006
+namespace CustomAvatar.Utilities
+{
+    internal struct AsyncInstantiateOperationAwaiter<T> : ICriticalNotifyCompletion where T : UnityEngine.Object
+    {
+        private AsyncInstantiateOperation<T> asyncOperation;
+        private Action continuationAction;
+
+        public AsyncInstantiateOperationAwaiter(AsyncInstantiateOperation<T> asyncOperation)
+        {
+            this.asyncOperation = asyncOperation;
+            this.continuationAction = null;
+        }
+
+        public bool IsCompleted => asyncOperation.isDone;
+
+        public AsyncInstantiateOperation<T> GetResult()
+        {
+            if (continuationAction != null)
+            {
+                asyncOperation.completed -= Continue;
+                continuationAction = null;
+            }
+
+            AsyncInstantiateOperation<T> op = asyncOperation;
+            asyncOperation = null;
+            return op;
+        }
+
+        public void OnCompleted(Action continuation)
+        {
+            UnsafeOnCompleted(continuation);
+        }
+
+        public void UnsafeOnCompleted(Action continuation)
+        {
+            if (continuationAction != null)
+            {
+                throw new InvalidOperationException("Continuation is already registered");
+            }
+
+            continuationAction = continuation;
+            asyncOperation.completed += Continue;
+        }
+
+        private void Continue(AsyncOperation _)
+        {
+            continuationAction?.Invoke();
+        }
+    }
+}
diff --git a/Source/CustomAvatar/Utilities/AsyncOperationExtensions.cs b/Source/CustomAvatar/Utilities/AsyncOperationExtensions.cs
index 44682ccf..225287d3 100644
--- a/Source/CustomAvatar/Utilities/AsyncOperationExtensions.cs
+++ b/Source/CustomAvatar/Utilities/AsyncOperationExtensions.cs
@@ -21,5 +21,7 @@ namespace CustomAvatar.Utilities
     internal static class AsyncOperationExtensions
     {
         public static AsyncOperationAwaiter<T> GetAwaiter<T>(this T asyncOperation) where T : AsyncOperation => new(asyncOperation);
+
+        public static AsyncInstantiateOperationAwaiter<T> GetAwaiter<T>(this AsyncInstantiateOperation<T> asyncOperation) where T : Object => new(asyncOperation);
     }
 }