Fixed null-reference exception when domain reload is enabled in Play Mode options

Author: SolidAlloy

Editor/Prefab Handling/ChangedPrefabs.cs

@@ -0,0 +1,116 @@
+namespace UnityHierarchyFolders.Editor
+{
+    using System;
+    using System.Collections;
+    using System.Collections.Generic;
+    using UnityEditor;
+    using UnityEngine;
+
+    /// <summary>
+    /// A singleton that contains info about edited prefabs and persists changes between domain reloads.
+    /// </summary>
+    [Serializable]
+    internal class ChangedPrefabs : IEnumerable<ValueTuple<string, string>>
+    {
+        private const string KeyName = nameof(ChangedPrefabs);
+
+        [SerializeField] private string[] _paths;
+        [SerializeField] private string[] _contents;
+
+        private static ChangedPrefabs _instance;
+
+        public static ChangedPrefabs Instance
+        {
+            get
+            {
+                // _instance is null only in PrefabFolderStripper.RevertChanges() when Instance is called for the first time.
+                // If _instance is null at that point, it means the domain reloaded, so the instance must be retrieved from PlayerPrefs.
+                // In all other cases, _instance is created with help of Initialize before operating on it, so FromDeserialized won't be called.
+                if (_instance == null)
+                {
+                    _instance = FromDeserialized();
+                }
+
+                return _instance;
+            }
+        }
+
+        public (string path, string content) this[int index]
+        {
+            get => (_paths[index], _contents[index]);
+            set
+            {
+                _paths[index] = value.path;
+                _contents[index] = value.content;
+            }
+        }
+
+        public static void Initialize(int length)
+        {
+            _instance = new ChangedPrefabs
+            {
+                _paths = new string[length],
+                _contents = new string[length]
+            };
+        }
+
+        public static void SerializeIfNeeded()
+        {
+            // Serialization is only needed if prefabs are edited before entering play mode and the domain will reload.
+            // In all other cases, changes to prefabs will be reverted before a domain reload.
+#if UNITY_2019_3_OR_NEWER
+            if (EditorSettings.enterPlayModeOptions.HasFlag(EnterPlayModeOptions.DisableDomainReload))
+                return;
+#endif
+
+            string serializedObject = EditorJsonUtility.ToJson(Instance);
+            PlayerPrefs.SetString(KeyName, serializedObject);
+        }
+
+        private static ChangedPrefabs FromDeserialized()
+        {
+            string serializedObject = PlayerPrefs.GetString(KeyName);
+            PlayerPrefs.DeleteKey(KeyName);
+            var instance = new ChangedPrefabs();
+            EditorJsonUtility.FromJsonOverwrite(serializedObject, instance);
+            return instance;
+        }
+
+        public Enumerator GetEnumerator() => new Enumerator(this);
+
+        IEnumerator<(string, string)> IEnumerable<ValueTuple<string, string>>.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+
+        IEnumerator IEnumerable.GetEnumerator()
+        {
+            return GetEnumerator();
+        }
+
+        public struct Enumerator : IEnumerator<ValueTuple<string, string>>
+        {
+            private readonly ChangedPrefabs _instance;
+            private int _index;
+
+            public Enumerator(ChangedPrefabs instance)
+            {
+                _instance = instance;
+                _index = -1;
+            }
+
+            public bool MoveNext()
+            {
+                return ++_index < Instance._paths.Length;
+            }
+
+            public void Reset() => _index = 0;
+
+            public (string, string) Current => _instance[_index];
+
+            object IEnumerator.Current => Current;
+
+            public void Dispose() { }
+        }
+    }
+}

Editor/Prefab Handling/ChangedPrefabs.cs.meta

@@ -8,13 +8,10 @@
     using UnityEditor.Build;
     using UnityEditor.Build.Reporting;
     using UnityEngine;
-    using Object = UnityEngine.Object;
 
     [InitializeOnLoad]
     public class PrefabFolderStripper : IPreprocessBuildWithReport, IPostprocessBuildWithReport
     {
-        private static (string path, string assetContent)[] _changedPrefabs;
-
         static PrefabFolderStripper()
         {
             EditorApplication.playModeStateChanged += HandlePrefabsOnPlayMode;
@@ -39,10 +36,12 @@ private static void HandlePrefabsOnPlayMode(PlayModeStateChange state)
             if ( ! StripSettings.StripFoldersFromPrefabsInPlayMode || StripSettings.PlayMode == StrippingMode.DoNothing)
                 return;
 
+            // Calling it not in EnteredPlayMode because scripts may instantiate prefabs in Awake or OnEnable
+            // which happens before EnteredPlayMode.
             if (state == PlayModeStateChange.ExitingEditMode)
             {
                 // Stripping folders from all prefabs in the project instead of only the ones referenced in the scenes
-                // because a prefab can be hot-swapped in Play Mode.
+                // because a prefab may be hot-swapped in Play Mode.
                 StripFoldersFromAllPrefabs();
             }
             else if (state == PlayModeStateChange.EnteredEditMode)
@@ -60,21 +59,23 @@ private static void StripFoldersFromDependentPrefabs()
                     AssetDatabase.GetLabels(GetAssetForLabel(path)).Contains(LabelHandler.FolderPrefabLabel))
                 .ToArray();
 
-            _changedPrefabs = new (string, string)[prefabsWithLabel.Length];
+            ChangedPrefabs.Initialize(prefabsWithLabel.Length);
 
             for (int i = 0; i < prefabsWithLabel.Length; i++)
             {
                 string path = prefabsWithLabel[i];
-                _changedPrefabs[i] = (path, File.ReadAllText(path));
+                ChangedPrefabs.Instance[i] = (path, File.ReadAllText(path));
                 StripFoldersFromPrefab(path, StripSettings.Build);
             }
+
+            // Serialization of ChangedPrefabs is not needed here because domain doesn't reload before changes are reverted.
         }
 
         private static
 #if UNITY_2020_1_OR_NEWER
             GUID
 #else
-            Object
+            UnityEngine.Object
 #endif
             GetAssetForLabel(string path)
         {
@@ -88,16 +89,20 @@ private static
         private static void StripFoldersFromAllPrefabs()
         {
             var prefabGUIDs = AssetDatabase.FindAssets($"l: {LabelHandler.FolderPrefabLabel}");
-            _changedPrefabs = new (string, string)[prefabGUIDs.Length];
+            ChangedPrefabs.Initialize(prefabGUIDs.Length);
 
             for (int i = 0; i < prefabGUIDs.Length; i++)
             {
                 string guid = prefabGUIDs[i];
                 string path = AssetDatabase.GUIDToAssetPath(guid);
 
-                _changedPrefabs[i] = (path, File.ReadAllText(path));
+                ChangedPrefabs.Instance[i] = (path, File.ReadAllText(path));
                 StripFoldersFromPrefab(path, StripSettings.PlayMode);
             }
+
+            // If domain reload is enabled in Play Mode Options, serialization of the changed prefabs is necessary
+            // so that changes can be reverted after leaving play mode.
+            ChangedPrefabs.SerializeIfNeeded();
         }
 
         private static void StripFoldersFromPrefab(string prefabPath, StrippingMode strippingMode)
@@ -115,7 +120,7 @@ private static void StripFoldersFromPrefab(string prefabPath, StrippingMode stri
 
         private static void RevertChanges()
         {
-            foreach ((string path, string content) in _changedPrefabs)
+            foreach ((string path, string content) in ChangedPrefabs.Instance)
             {
                 File.WriteAllText(path, content);
             }