Merge pull request #1219 from anatawa12/fix-prefab-safe-set-api-with-…

…on-scene-layer

fix: prefab safe set API is broken with onSceneLayer

CHANGELOG-PRERELEASE.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog].
 ### Removed
 
 ### Fixed
+- API about Prefab Safe Set are broken with prefab instance `#1219`
 
 ### Security
 

CHANGELOG.md

@@ -46,7 +46,7 @@ The format is based on [Keep a Changelog].
 - Reimplement Preview system with NDMF Preview System `#1131` `#1195` `#1218`
   - This will prevent issues relates to Animation Mode bug.
   - This allows you to preview Remove Mesh components without selecting Mesh OR while in Animation Mode.
-- Improved Prefab Safe Set, which are used in MergePhysBone, MergeSkinnedMesh, FreezeBlendShape and more components `#1212`
+- Improved Prefab Safe Set, which are used in MergePhysBone, MergeSkinnedMesh, FreezeBlendShape and more components `#1212` `#1219`
   - This should improve compatibility with replacing base prefab, which is added in Unity 2022.
 - Allow multiple component for Remove Mesh components with API `#1216` `#1218`
   - This allows non-destructive tools to add Remove Mesh components even if Remove Mesh component are added before.

Internal/PrefabSafeSet/Editor/Utils.cs

@@ -20,20 +20,8 @@ public static int PrefabNestCount(Object instance)
             return nestCount;
         }
 
-        public static bool ShouldUsePrefabOnSceneLayer(Object instance)
-        {
-            var isInstance = PrefabUtility.IsPartOfPrefabInstance(instance);
-            var isAsset = PrefabUtility.IsPartOfPrefabAsset(instance);
-
-            var currentPrefabStage = PrefabStageUtility.GetCurrentPrefabStage();
-            if (currentPrefabStage != null)
-            {
-                var instanceGameObject = instance as GameObject ?? (instance as Component)?.gameObject;
-                isAsset |= currentPrefabStage.IsPartOfPrefabContents(instanceGameObject);
-            }
-
-            return isInstance && !isAsset;
-        }
+        public static bool ShouldUsePrefabOnSceneLayer(Object instance) =>
+            PrefabSafeSetRuntimeUtil.ShouldUsePrefabOnSceneLayer(instance);
 
         public static bool IsNullOrMissing<T>(this T self, Object context) =>
             self.IsNullOrMissing(new NullOrMissingContext(context));

Internal/PrefabSafeSet/Runtime/PrefabSafeSet.cs

@@ -10,6 +10,23 @@ namespace Anatawa12.AvatarOptimizer.PrefabSafeSet
 {
     internal static class PrefabSafeSetRuntimeUtil
     {
+#if UNITY_EDITOR
+        public static bool ShouldUsePrefabOnSceneLayer(Object instance)
+        {
+            var isInstance = UnityEditor.PrefabUtility.IsPartOfPrefabInstance(instance);
+            var isAsset = UnityEditor.PrefabUtility.IsPartOfPrefabAsset(instance);
+
+            var currentPrefabStage = UnityEditor.SceneManagement.PrefabStageUtility.GetCurrentPrefabStage();
+            if (currentPrefabStage != null)
+            {
+                var instanceGameObject = instance as GameObject ?? (instance as Component)?.gameObject;
+                isAsset |= currentPrefabStage.IsPartOfPrefabContents(instanceGameObject);
+            }
+
+            return isInstance && !isAsset;
+        }
+#endif
+
         public static void ResizeArray<T>(ref T[] array, int size) where T : new()
         {
             var source = array;
@@ -123,6 +140,7 @@ public override HashSet<T> GetAsSet()
             var result = new HashSet<T>(mainSet.Where(x => x.IsNotNull()));
             foreach (var layer in prefabLayers)
                 layer.ApplyTo(result);
+            onSceneLayer.ApplyTo(result);
             return result;
         }
 
@@ -132,19 +150,32 @@ public override List<T> GetAsList()
             var set = new HashSet<T>(result);
             foreach (var layer in prefabLayers)
                 layer.ApplyTo(set, result);
+            onSceneLayer.ApplyTo(set, result);
             return result;
         }
 
 #if UNITY_EDITOR
-        private (HashSet<T>, PrefabLayer<T>) GetBaseSetAndLayer(int nestCount)
+        private (HashSet<T>, PrefabLayer<T>) GetBaseSetAndLayer(int nestCount, bool useOnSceneLayer)
         {
-            if (prefabLayers.Length < nestCount)
-                PrefabSafeSetRuntimeUtil.ResizeArray(ref prefabLayers, nestCount);
-            var baseSet = new HashSet<T>(mainSet.Where(x => x.IsNotNull()));
-            for (var i = 0; i < nestCount - 1; i++) prefabLayers[i].ApplyTo(baseSet);
-            var layer = prefabLayers[nestCount - 1];
+            if (useOnSceneLayer)
+            {
+                if (!usingOnSceneLayer)
+                    usingOnSceneLayer = true;
+                var baseSet = new HashSet<T>(mainSet.Where(x => x.IsNotNull()));
+                for (var i = 0; i < prefabLayers.Length && i < nestCount - 1; i++) prefabLayers[i].ApplyTo(baseSet);
 
-            return (baseSet, layer);
+                return (baseSet, onSceneLayer);
+            }
+            else
+            {
+                if (prefabLayers.Length < nestCount)
+                    PrefabSafeSetRuntimeUtil.ResizeArray(ref prefabLayers, nestCount);
+                var baseSet = new HashSet<T>(mainSet.Where(x => x.IsNotNull()));
+                for (var i = 0; i < nestCount - 1; i++) prefabLayers[i].ApplyTo(baseSet);
+                var layer = prefabLayers[nestCount - 1];
+
+                return (baseSet, layer);
+            }
         }
 
         private static int PrefabNestCount(Object instance)
@@ -160,6 +191,7 @@ public override bool AddRange(IEnumerable<T> values)
         {
             var valueEnumerable = values.Where(x => x.IsNotNull());
             var nestCount = PrefabNestCount(OuterObject);
+            var useOnSceneLayer = PrefabSafeSetRuntimeUtil.ShouldUsePrefabOnSceneLayer(OuterObject);
 
             if (nestCount == 0)
             {
@@ -169,7 +201,7 @@ public override bool AddRange(IEnumerable<T> values)
             }
             else
             {
-                var (baseSet, layer) = GetBaseSetAndLayer(nestCount);
+                var (baseSet, layer) = GetBaseSetAndLayer(nestCount, useOnSceneLayer);
                 var valuesList = new List<T>(valueEnumerable);
 
                 var originalRemoves = layer.removes.Length;
@@ -186,6 +218,7 @@ public override bool RemoveRange(IEnumerable<T> values)
         {
             var valueEnumerable = values.Where(x => x.IsNotNull());
             var nestCount = PrefabNestCount(OuterObject);
+            var useOnSceneLayer = PrefabSafeSetRuntimeUtil.ShouldUsePrefabOnSceneLayer(OuterObject);
 
             if (nestCount == 0)
             {
@@ -195,7 +228,7 @@ public override bool RemoveRange(IEnumerable<T> values)
             }
             else
             {
-                var (baseSet, layer) = GetBaseSetAndLayer(nestCount);
+                var (baseSet, layer) = GetBaseSetAndLayer(nestCount, useOnSceneLayer);
                 var valuesList = new List<T>(valueEnumerable);
 
                 var originalRemoves = layer.removes.Length;
@@ -211,14 +244,15 @@ public override bool RemoveRange(IEnumerable<T> values)
         public override void RemoveIf(Func<T, bool> predicate)
         {
             var nestCount = PrefabNestCount(OuterObject);
+            var useOnSceneLayer = PrefabSafeSetRuntimeUtil.ShouldUsePrefabOnSceneLayer(OuterObject);
 
             if (nestCount == 0)
             {
                 mainSet = mainSet.Where(x => x.IsNull() || !predicate(x)).ToArray();
             }
             else
             {
-                var (baseSet, layer) = GetBaseSetAndLayer(nestCount);
+                var (baseSet, layer) = GetBaseSetAndLayer(nestCount, useOnSceneLayer);
 
                 layer.removes = layer.removes.Concat(baseSet.Where(predicate)).ToArray();
                 layer.additions = layer.additions.Where(x => !predicate(x)).ToArray();
@@ -228,14 +262,15 @@ public override void RemoveIf(Func<T, bool> predicate)
         public override void Clear()
         {
             var nestCount = PrefabNestCount(OuterObject);
+            var useSceneLayer = PrefabSafeSetRuntimeUtil.ShouldUsePrefabOnSceneLayer(OuterObject);
 
             if (nestCount == 0)
             {
                 mainSet = Array.Empty<T>();
             }
             else
             {
-                var (baseSet, layer) = GetBaseSetAndLayer(nestCount);
+                var (baseSet, layer) = GetBaseSetAndLayer(nestCount, useSceneLayer);
 
                 layer.removes = layer.removes.Concat(baseSet).ToArray();
                 layer.additions = Array.Empty<T>();