From 5c0b62f295eac9197711479163b7650ecbb199bc Mon Sep 17 00:00:00 2001
From: bd_ <bdunderscore@fushizen.net>
Date: Sat, 10 Sep 2022 10:47:41 -0700
Subject: [PATCH] Restore Merge Armature locked mode references on prefab
 instantiation

---
 .../Editor/Util.cs                            |  8 ++++-
 .../Runtime/ModularAvatarMergeArmature.cs     | 33 ++++++++++++-------
 .../Runtime/RuntimeUtil.cs                    |  5 +++
 3 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/Packages/net.fushizen.modular-avatar/Editor/Util.cs b/Packages/net.fushizen.modular-avatar/Editor/Util.cs
index 0625b13e..ff6813f5 100644
--- a/Packages/net.fushizen.modular-avatar/Editor/Util.cs
+++ b/Packages/net.fushizen.modular-avatar/Editor/Util.cs
@@ -37,11 +37,17 @@ namespace net.fushizen.modular_avatar.core.editor
         }
     }
     
+    [InitializeOnLoad]
     public static class Util
     {
         internal const string generatedAssetsSubdirectory = "999_Modular_Avatar_Generated";
         internal const string generatedAssetsPath = "Assets/" + generatedAssetsSubdirectory;
 
+        static Util()
+        {
+            RuntimeUtil.delayCall = (cb) => EditorApplication.delayCall += cb.Invoke;
+        }
+
         static internal AnimatorController CreateContainer()
         {
             var container = new AnimatorController();
@@ -74,7 +80,7 @@ namespace net.fushizen.modular_avatar.core.editor
                 var subdir = generatedAssetsPath;
 
                 AssetDatabase.DeleteAsset(subdir);
-                //FileUtil.DeleteFileOrDirectory(subdir);
+                FileUtil.DeleteFileOrDirectory(subdir);
             };
         }
     }
diff --git a/Packages/net.fushizen.modular-avatar/Runtime/ModularAvatarMergeArmature.cs b/Packages/net.fushizen.modular-avatar/Runtime/ModularAvatarMergeArmature.cs
index 04348b8b..cb92d220 100644
--- a/Packages/net.fushizen.modular-avatar/Runtime/ModularAvatarMergeArmature.cs
+++ b/Packages/net.fushizen.modular-avatar/Runtime/ModularAvatarMergeArmature.cs
@@ -23,9 +23,6 @@
  */
 
 using System;
-#if UNITY_EDITOR
-using UnityEditor;
-#endif
 using UnityEngine;
 
 namespace net.fushizen.modular_avatar.core
@@ -40,10 +37,9 @@ namespace net.fushizen.modular_avatar.core
         public bool locked;
 
         private bool wasLocked;
-#if UNITY_EDITOR
         void OnValidate()
         {
-            EditorApplication.delayCall += () =>
+            RuntimeUtil.delayCall(() =>
             {
                 if (this == null) return;
                 if (mergeTarget == null && !string.IsNullOrWhiteSpace(mergeTargetPath))
@@ -53,7 +49,9 @@ namespace net.fushizen.modular_avatar.core
                     {
                         mergeTarget = avatar.transform.Find(mergeTargetPath)?.gameObject;
                     }
-                    if (mergeTarget != null) {
+
+                    if (mergeTarget != null)
+                    {
                         RuntimeUtil.MarkDirty(this);
                     }
                 }
@@ -69,9 +67,8 @@ namespace net.fushizen.modular_avatar.core
                 }
 
                 CheckLock();
-            };
+            });
         }
-#endif
 
         void CheckLock()
         {
@@ -94,11 +91,23 @@ namespace net.fushizen.modular_avatar.core
                     foreach (var xform in GetComponentsInChildren<Transform>(true))
                     {
                         Transform baseObject = FindCorresponding(xform);
-                        if (baseObject != null && xform.gameObject.GetComponent<MAInternalOffsetMarker>() == null)
+                        var marker = xform.gameObject.GetComponent<MAInternalOffsetMarker>();
+
+                        if (baseObject == null)
                         {
-                            var comp = xform.gameObject.AddComponent<MAInternalOffsetMarker>();
-                            comp.correspondingObject = baseObject;
-                            comp.lockBasePosition = baseObject.gameObject == mergeTarget;
+                            if (marker != null)
+                            {
+                                DestroyImmediate(marker);
+                            }
+                        }
+                        else
+                        {
+                            if (marker == null)
+                            {
+                                marker = xform.gameObject.AddComponent<MAInternalOffsetMarker>();
+                            }
+                            marker.correspondingObject = baseObject;
+                            marker.lockBasePosition = baseObject.gameObject == mergeTarget;
                         }
                     }
 
diff --git a/Packages/net.fushizen.modular-avatar/Runtime/RuntimeUtil.cs b/Packages/net.fushizen.modular-avatar/Runtime/RuntimeUtil.cs
index 7240f693..3c3d24c1 100644
--- a/Packages/net.fushizen.modular-avatar/Runtime/RuntimeUtil.cs
+++ b/Packages/net.fushizen.modular-avatar/Runtime/RuntimeUtil.cs
@@ -33,6 +33,11 @@ namespace net.fushizen.modular_avatar.core
 {
     public static class RuntimeUtil
     {
+        public delegate void NullCallback();
+
+        // Initialized in Util
+        public static Action<NullCallback> delayCall = (_) => { };
+        
         [CanBeNull]
         public static string RelativePath(GameObject root, GameObject child)
         {