Add "one-shot" option for TransformTree updating

Author: sinai-dev

src/Inspectors/GameObjectInspector.cs

@@ -82,7 +82,7 @@ public void ChangeTarget(GameObject newTarget)
             this.Target = newTarget;
             GOControls.UpdateGameObjectInfo(true, true);
             GOControls.UpdateTransformControlValues(true);
-            TransformTree.RefreshData(true, false, true);
+            TransformTree.RefreshData(true, false, true, false);
             UpdateComponents();
         }
 
@@ -109,7 +109,7 @@ public override void Update()
 
                 GOControls.UpdateGameObjectInfo(false, false);
 
-                TransformTree.RefreshData(true, false, false);
+                TransformTree.RefreshData(true, false, false, false);
                 UpdateComponents();
             }
         }
@@ -220,7 +220,7 @@ private void OnAddChildClicked(string input)
             var newObject = new GameObject(input);
             newObject.transform.parent = GOTarget.transform;
 
-            TransformTree.RefreshData(true, false, true);
+            TransformTree.RefreshData(true, false, true, false);
         }
 
         private void OnAddComponentClicked(string input)

src/ObjectExplorer/SceneExplorer.cs

@@ -64,7 +64,7 @@ public void Update()
         public void UpdateTree()
         {
             SceneHandler.Update();
-            Tree.RefreshData(true, false, false);
+            Tree.RefreshData(true, false, false, false);
         }
 
         public void JumpToTransform(Transform transform)
@@ -94,7 +94,7 @@ private void OnSceneSelectionDropdownChanged(int value)
 
             SceneHandler.SelectedScene = SceneHandler.LoadedScenes[value];
             SceneHandler.Update();
-            Tree.RefreshData(true, true, true);
+            Tree.RefreshData(true, true, true, false);
             OnSelectedSceneChanged(SceneHandler.SelectedScene.Value);
         }
 
@@ -158,7 +158,7 @@ private void OnFilterInput(string input)
             }
 
             Tree.CurrentFilter = input;
-            Tree.RefreshData(true, false, true);
+            Tree.RefreshData(true, false, true, false);
         }
 
         private void TryLoadScene(LoadSceneMode mode, Dropdown allSceneDrop)
@@ -259,7 +259,7 @@ public override void ConstructUI(GameObject content)
 
             Tree = new TransformTree(scrollPool, GetRootEntries);
             Tree.Init();
-            Tree.RefreshData(true, true, true);
+            Tree.RefreshData(true, true, true, false);
             //scrollPool.Viewport.GetComponent<Mask>().enabled = false;
             //UIRoot.GetComponent<Mask>().enabled = false;
 

src/UI/Widgets/TransformTree/CachedTransform.cs

@@ -19,7 +19,7 @@ public class CachedTransform
         public bool Enabled { get; internal set; }
         public int SiblingIndex { get; internal set; }
 
-        public bool Expanded => Tree.IsCellExpanded(InstanceID);
+        public bool Expanded => Tree.IsTransformExpanded(InstanceID);
 
         public CachedTransform(TransformTree tree, Transform transform, int depth, CachedTransform parent = null)
         {

src/UI/Widgets/TransformTree/TransformTree.cs

@@ -17,6 +17,13 @@ public class TransformTree : ICellPoolDataSource<TransformCell>
 
         public ScrollPool<TransformCell> ScrollPool;
 
+        // IMPORTANT CAVEAT WITH OrderedDictionary:
+        // While the performance is mostly good, there are two methods we should NEVER use:
+        // - Remove(object)
+        // - set_Item[object]
+        // These two methods have extremely bad performance due to using IndexOfKey(), which iterates the whole dictionary.
+        // Currently we do not use either of these methods, so everything should be constant time hash lookups.
+        // We DO make use of get_Item[object], get_Item[index], Add, Insert and RemoveAt, which OrderedDictionary perfectly meets our needs for.
         /// <summary>
         /// Key: UnityEngine.Transform instance ID<br/>
         /// Value: CachedTransform
@@ -27,13 +34,19 @@ public class TransformTree : ICellPoolDataSource<TransformCell>
         private readonly HashSet<int> expandedInstanceIDs = new();
         private readonly HashSet<int> autoExpandedIDs = new();
 
+        // state for Traverse parse
         private readonly HashSet<int> visited = new();
-        private bool needRefresh;
+        private bool needRefreshUI;
         private int displayIndex;
         int prevDisplayIndex;
 
+        private Coroutine refreshCoroutine;
+        private readonly Stopwatch traversedThisFrame = new();
+
+        // ScrollPool item count. PrevDisplayIndex is the highest index + 1 from our last traverse.
         public int ItemCount => prevDisplayIndex;
 
+        // Search filter
         public bool Filtering => !string.IsNullOrEmpty(currentFilter);
         private bool wasFiltering;
 
@@ -54,35 +67,52 @@ public string CurrentFilter
         }
         private string currentFilter;
 
-        private Coroutine refreshCoroutine;
-        private readonly Stopwatch traversedThisFrame = new();
-
         public TransformTree(ScrollPool<TransformCell> scrollPool, Func<IEnumerable<GameObject>> getRootEntriesMethod)
         {
             ScrollPool = scrollPool;
             GetRootEntriesMethod = getRootEntriesMethod;
         }
 
+        // Initialize and reset
+
+        // Must be called externally from owner of this TransformTree
         public void Init()
         {
             ScrollPool.Initialize(this);
         }
 
+        // Called to completely reset the tree, ie. switching inspected GameObject
+        public void Rebuild()
+        {
+            autoExpandedIDs.Clear();
+            expandedInstanceIDs.Clear();
+
+            RefreshData(true, true, true, false);
+        }
+
+        // Called to completely wipe our data (ie, GameObject inspector returning to pool)
         public void Clear()
         {
             this.cachedTransforms.Clear();
             displayIndex = 0;
             autoExpandedIDs.Clear();
             expandedInstanceIDs.Clear();
             this.ScrollPool.Refresh(true, true);
+            if (refreshCoroutine != null)
+            {
+                RuntimeHelper.StopCoroutine(refreshCoroutine);
+                refreshCoroutine = null;
+            }
         }
 
-        public bool IsCellExpanded(int instanceID)
+        // Checks if the given Instance ID is expanded or not
+        public bool IsTransformExpanded(int instanceID)
         {
             return Filtering ? autoExpandedIDs.Contains(instanceID)
                              : expandedInstanceIDs.Contains(instanceID);
         }
 
+        // Jumps to a specific Transform in the tree and highlights it.
         public void JumpAndExpandToTransform(Transform transform)
         {
             // make sure all parents of the object are expanded
@@ -96,8 +126,9 @@ public void JumpAndExpandToTransform(Transform transform)
                 parent = parent.parent;
             }
 
-            // Refresh cached transforms (no UI rebuild yet)
-            RefreshData(false, false, false);
+            // Refresh cached transforms (no UI rebuild yet).
+            // Stop existing coroutine and do it oneshot.
+            RefreshData(false, false, true, true);
 
             int transformID = transform.GetInstanceID();
 
@@ -130,15 +161,9 @@ private IEnumerator HighlightCellCoroutine(TransformCell cell)
             button.OnDeselect(null);
         }
 
-        public void Rebuild()
-        {
-            autoExpandedIDs.Clear();
-            expandedInstanceIDs.Clear();
-
-            RefreshData(true, true, true);
-        }
-
-        public void RefreshData(bool andRefreshUI, bool jumpToTop, bool stopExistingCoroutine)
+        // Perform a Traverse and optionally refresh the ScrollPool as well.
+        // If oneShot, then this happens instantly with no yield.
+        public void RefreshData(bool andRefreshUI, bool jumpToTop, bool stopExistingCoroutine, bool oneShot)
         {
             if (refreshCoroutine != null)
             {
@@ -153,25 +178,29 @@ public void RefreshData(bool andRefreshUI, bool jumpToTop, bool stopExistingCoro
 
             visited.Clear();
             displayIndex = 0;
-            needRefresh = false;
+            needRefreshUI = false;
             traversedThisFrame.Reset();
             traversedThisFrame.Start();
 
             IEnumerable<GameObject> rootObjects = GetRootEntriesMethod.Invoke();
 
-            refreshCoroutine = RuntimeHelper.StartCoroutine(RefreshCoroutine(rootObjects, andRefreshUI, jumpToTop));
+            refreshCoroutine = RuntimeHelper.StartCoroutine(RefreshCoroutine(rootObjects, andRefreshUI, jumpToTop, oneShot));
         }
 
-        private IEnumerator RefreshCoroutine(IEnumerable<GameObject> rootObjects, bool andRefreshUI, bool jumpToTop)
+        // Coroutine for batched updates, max 2000 gameobjects per frame so FPS doesn't get tanked when there is like 100k gameobjects.
+        // if "oneShot", then this will NOT be batched (if we need an immediate full update).
+        IEnumerator RefreshCoroutine(IEnumerable<GameObject> rootObjects, bool andRefreshUI, bool jumpToTop, bool oneShot)
         {
-            var thisCoro = refreshCoroutine;
             foreach (var gameObj in rootObjects)
             {
                 if (gameObj)
                 {
-                    var enumerator = Traverse(gameObj.transform);
+                    var enumerator = Traverse(gameObj.transform, null, 0, oneShot);
                     while (enumerator.MoveNext())
-                        yield return enumerator.Current;
+                    {
+                        if (!oneShot)
+                            yield return enumerator.Current;
+                    }
                 }
             }
 
@@ -182,17 +211,20 @@ private IEnumerator RefreshCoroutine(IEnumerable<GameObject> rootObjects, bool a
                 if (!visited.Contains(cached.InstanceID))
                 {
                     cachedTransforms.RemoveAt(i);
-                    needRefresh = true;
+                    needRefreshUI = true;
                 }
             }
 
-            if (andRefreshUI && needRefresh)
+            if (andRefreshUI && needRefreshUI)
                 ScrollPool.Refresh(true, jumpToTop);
 
             prevDisplayIndex = displayIndex;
-        }
+            refreshCoroutine = null;
+        }   
 
-        private IEnumerator Traverse(Transform transform, CachedTransform parent = null, int depth = 0)
+        // Recursive method to check a Transform and its children (if expanded).
+        // Parent and depth can be null/default.
+        private IEnumerator Traverse(Transform transform, CachedTransform parent, int depth, bool oneShot)
         {
             // Let's only tank 2ms of each frame (60->53fps)
             if (traversedThisFrame.ElapsedMilliseconds > 2)
@@ -227,7 +259,7 @@ private IEnumerator Traverse(Transform transform, CachedTransform parent = null,
                 int prevSiblingIdx = cached.SiblingIndex;
                 if (cached.Update(transform, depth))
                 {
-                    needRefresh = true;
+                    needRefreshUI = true;
 
                     // If the sibling index changed, we need to shuffle it in our cached transforms list.
                     if (prevSiblingIdx != cached.SiblingIndex)
@@ -242,7 +274,7 @@ private IEnumerator Traverse(Transform transform, CachedTransform parent = null,
             }
             else
             {
-                needRefresh = true;
+                needRefreshUI = true;
                 cached = new CachedTransform(this, transform, depth, parent);
                 if (cachedTransforms.Count <= displayIndex)
                     cachedTransforms.Add(instanceID, cached);
@@ -252,13 +284,16 @@ private IEnumerator Traverse(Transform transform, CachedTransform parent = null,
 
             displayIndex++;
 
-            if (IsCellExpanded(instanceID) && cached.Value.childCount > 0)
+            if (IsTransformExpanded(instanceID) && cached.Value.childCount > 0)
             {
                 for (int i = 0; i < transform.childCount; i++)
                 {
-                    var enumerator = Traverse(transform.GetChild(i), cached, depth + 1);
+                    var enumerator = Traverse(transform.GetChild(i), cached, depth + 1, oneShot);
                     while (enumerator.MoveNext())
-                        yield return enumerator.Current;
+                    {
+                        if (!oneShot)
+                            yield return enumerator.Current;
+                    }
                 }
             }
         }
@@ -316,14 +351,14 @@ public void OnCellExpandToggled(CachedTransform cache)
             else
                 expandedInstanceIDs.Add(instanceID);
 
-            RefreshData(true, false, true);
+            RefreshData(true, false, true, false);
         }
 
         public void OnCellEnableToggled(CachedTransform cache)
         {
             cache.Value.gameObject.SetActive(!cache.Value.gameObject.activeSelf);
 
-            RefreshData(true, false, true);
+            RefreshData(true, false, true, false);
         }
     }
 }