perf: avoid loading assets unnecessarily from ChangeNotifier (#487)

Reported in #1357

CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 ### Fixed
+- [#487] Fixed a performance issue where all assets would potentially be loaded on reimport, taking a lot of time and
+  memory in the process
 
 ### Changed
 

Editor/ChangeStream/ChangeNotifier.cs

@@ -1,5 +1,6 @@
 #region
 
+using System.Collections.Generic;
 using JetBrains.Annotations;
 using nadena.dev.ndmf.cs;
 using UnityEditor;
@@ -12,6 +13,22 @@ namespace nadena.dev.ndmf.preview
     [PublicAPI]
     public static class ChangeNotifier
     {
+        private static Dictionary<string, HashSet<int>> pathToInstanceIds = new();
+
+        internal static void RecordObjectOfInterest(Object obj)
+        {
+            if (!AssetDatabase.Contains(obj)) return;
+
+            var path = AssetDatabase.GetAssetPath(obj);
+            if (!pathToInstanceIds.TryGetValue(path, out var set))
+            {
+                set = new HashSet<int>();
+                pathToInstanceIds[path] = set;
+            }
+
+            set.Add(obj.GetInstanceID());
+        }
+
         /// <summary>
         /// Notifies the reactive query and NDMF preview system of a change in an object that isn't tracked by the normal
         /// unity ObjectchangeEventStream system.
@@ -26,12 +43,23 @@ public static void NotifyObjectUpdate(Object obj)
         /// Notifies the reactive query and NDMF preview system of a change in an object that isn't tracked by the normal
         /// unity ObjectchangeEventStream system.
         /// </summary>
-        /// <param name="obj"></param>
+        /// <param name="instanceId"></param>
         public static void NotifyObjectUpdate(int instanceId)
         {
             ObjectWatcher.Instance.Hierarchy.InvalidateTree(instanceId);
         }
 
+        private static void NotifyAssetFileChange(string path)
+        {
+            if (pathToInstanceIds.TryGetValue(path, out var set))
+            {
+                foreach (var instanceId in set)
+                {
+                    NotifyObjectUpdate(instanceId);
+                }
+            }
+        }
+
         private class Processor : AssetPostprocessor
         {
             [UsedImplicitly]
@@ -48,14 +76,16 @@ static void OnPostprocessAllAssets(
                 bool didDomainReload
             )
             {
-                foreach (var asset in importedAssets)
+                using var _ = ObjectWatcher.Instance.Hierarchy.SuspendEvents();
+
+                foreach (var path in importedAssets)
+                {
+                    NotifyAssetFileChange(path);
+                }
+
+                foreach (var path in deletedAssets)
                 {
-                    if (asset.EndsWith(".unity")) continue;
-                    var subassets = AssetDatabase.LoadAllAssetsAtPath(asset);
-                    foreach (var subasset in subassets)
-                    {
-                        NotifyObjectUpdate(subasset);
-                    }
+                    NotifyAssetFileChange(path);
                 }
             }
         }

Editor/ChangeStream/ShadowGameObject.cs

@@ -318,23 +318,25 @@ ComputeContext ctx
             return shadowObject._listeners.Register(filter, ctx);
         }
 
-        internal IDisposable RegisterObjectListener(UnityObject targetComponent,
+        internal IDisposable RegisterObjectListener(UnityObject targetObject,
             ListenerSet<HierarchyEvent>.Filter filter,
             ComputeContext ctx
         )
         {
-            if (targetComponent == null || ctx.IsInvalidated) return new NullDisposable();
+            if (targetObject == null || ctx.IsInvalidated) return new NullDisposable();
 
 #if NDMF_TRACE_SHADOW
             System.Diagnostics.Debug.WriteLine($"[ShadowHierarchy] RegisterObjectListener({targetComponent.GetInstanceID()})");
 #endif
 
-            if (!_otherObjects.TryGetValue(targetComponent.GetInstanceID(), out var shadowComponent))
+            if (!_otherObjects.TryGetValue(targetObject.GetInstanceID(), out var shadowComponent))
             {
-                shadowComponent = new ShadowObject(targetComponent);
-                _otherObjects[targetComponent.GetInstanceID()] = shadowComponent;
+                shadowComponent = new ShadowObject(targetObject);
+                _otherObjects[targetObject.GetInstanceID()] = shadowComponent;
             }
 
+            ChangeNotifier.RecordObjectOfInterest(targetObject);
+
             return shadowComponent._listeners.Register(filter, ctx);
         }
 

UnitTests~/PreviewSystem/ChangeNotifierTest.cs

@@ -0,0 +1,48 @@
+using System.Collections;
+using System.Collections.Generic;
+using nadena.dev.ndmf.preview;
+using NUnit.Framework;
+using UnityEditor;
+using UnityEngine;
+using UnityEngine.TestTools;
+
+namespace UnitTests
+{
+    public class ChangeNotifierTest : TestBase
+    {
+        [Test]
+        public void WhenAssetReimported_InvalidatesListeners()
+        {
+            var path = "Assets/ChangeNotifierTest.txt";
+            using (var f = System.IO.File.CreateText(path))
+            {
+                f.WriteLine("Hello, world!");
+            }
+
+            AssetDatabase.Refresh();
+
+            var asset = AssetDatabase.LoadAssetAtPath<TextAsset>(path);
+
+            var ctx = new ComputeContext("test");
+            ctx.Observe(asset);
+
+            using (var f = System.IO.File.CreateText(path))
+            {
+                f.WriteLine("Goodbye, world!");
+            }
+
+            Assert.IsFalse(ctx.IsInvalidated);
+
+            AssetDatabase.Refresh();
+
+            Assert.IsTrue(ctx.IsInvalidated);
+
+            ctx = new ComputeContext("test");
+            ctx.Observe(asset);
+
+            AssetDatabase.DeleteAsset(path);
+
+            Assert.IsTrue(ctx.IsInvalidated);
+        }
+    }
+}