Improve TransformTree efficiency

- Added batching to the update method so that a maximum of 2000 GameObjects are traversed each frame.
- Changed from OrderedDictionary.Remove to OrderedDictionary.RemoveAt when pruning entries as the former needs to iterate through all entries to find the index of the key, whereas the latter is constant time.

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);
+            TransformTree.RefreshData(true, false, true);
             UpdateComponents();
         }
 
@@ -109,7 +109,7 @@ public override void Update()
 
                 GOControls.UpdateGameObjectInfo(false, false);
 
-                TransformTree.RefreshData(true, false);
+                TransformTree.RefreshData(true, 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);
+            TransformTree.RefreshData(true, false, true);
         }
 
         private void OnAddComponentClicked(string input)

src/ObjectExplorer/SceneExplorer.cs

@@ -64,7 +64,7 @@ public void Update()
         public void UpdateTree()
         {
             SceneHandler.Update();
-            Tree.RefreshData(true);
+            Tree.RefreshData(true, 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);
+            Tree.RefreshData(true, true, true);
             OnSelectedSceneChanged(SceneHandler.SelectedScene.Value);
         }
 
@@ -158,7 +158,7 @@ private void OnFilterInput(string input)
             }
 
             Tree.CurrentFilter = input;
-            Tree.RefreshData(true, true);
+            Tree.RefreshData(true, false, true);
         }
 
         private void TryLoadScene(LoadSceneMode mode, Dropdown allSceneDrop)
@@ -239,6 +239,17 @@ public override void ConstructUI(GameObject content)
 
             refreshRow.SetActive(false);
 
+            // tree labels row
+
+            var labelsRow = UIFactory.CreateHorizontalGroup(toolbar, "LabelsRow", true, true, true, true, 2, new Vector4(2, 2, 2, 2));
+            UIFactory.SetLayoutElement(labelsRow, minHeight: 30, flexibleHeight: 0);
+
+            var nameLabel = UIFactory.CreateLabel(labelsRow, "NameLabel", "Name", TextAnchor.MiddleLeft, color: Color.grey);
+            UIFactory.SetLayoutElement(nameLabel.gameObject, flexibleWidth: 9999, minHeight: 25);
+
+            var indexLabel = UIFactory.CreateLabel(labelsRow, "IndexLabel", "Sibling Index", TextAnchor.MiddleLeft, fontSize: 12, color: Color.grey);
+            UIFactory.SetLayoutElement(indexLabel.gameObject, minWidth: 100, flexibleWidth: 0, minHeight: 25);
+
             // Transform Tree
 
             var scrollPool = UIFactory.CreateScrollPool<TransformCell>(uiRoot, "TransformTree", out GameObject scrollObj,
@@ -248,7 +259,7 @@ public override void ConstructUI(GameObject content)
 
             Tree = new TransformTree(scrollPool, GetRootEntries);
             Tree.Init();
-            Tree.RefreshData(true, true);
+            Tree.RefreshData(true, true, true);
             //scrollPool.Viewport.GetComponent<Mask>().enabled = false;
             //UIRoot.GetComponent<Mask>().enabled = false;
 

src/UI/Widgets/TransformTree/TransformTree.cs

@@ -2,12 +2,9 @@
 using System.Collections;
 using System.Collections.Generic;
 using System.Collections.Specialized;
-using System.Linq;
-using System.Text;
+using System.Diagnostics;
 using UnityEngine;
-using UnityEngine.UI;
 using UniverseLib;
-using UniverseLib.UI.Widgets;
 using UniverseLib.UI.Widgets.ScrollView;
 using UniverseLib.Utility;
 
@@ -24,17 +21,18 @@ public class TransformTree : ICellPoolDataSource<TransformCell>
         /// Key: UnityEngine.Transform instance ID<br/>
         /// Value: CachedTransform
         /// </summary>
-        internal readonly OrderedDictionary cachedTransforms = new OrderedDictionary();
+        internal readonly OrderedDictionary cachedTransforms = new();
 
         // for keeping track of which actual transforms are expanded or not, outside of the cache data.
-        private readonly HashSet<int> expandedInstanceIDs = new HashSet<int>();
-        private readonly HashSet<int> autoExpandedIDs = new HashSet<int>();
+        private readonly HashSet<int> expandedInstanceIDs = new();
+        private readonly HashSet<int> autoExpandedIDs = new();
 
-        private readonly HashSet<int> visited = new HashSet<int>();
+        private readonly HashSet<int> visited = new();
         private bool needRefresh;
         private int displayIndex;
+        int prevDisplayIndex;
 
-        public int ItemCount => cachedTransforms.Count;
+        public int ItemCount => prevDisplayIndex;
 
         public bool Filtering => !string.IsNullOrEmpty(currentFilter);
         private bool wasFiltering;
@@ -56,45 +54,15 @@ 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;
         }
 
-        public void OnCellBorrowed(TransformCell cell)
-        {
-            cell.OnExpandToggled += OnCellExpandToggled;
-            cell.OnGameObjectClicked += OnGameObjectClicked;
-            cell.OnEnableToggled += OnCellEnableToggled;
-        }
-
-        private void OnGameObjectClicked(GameObject obj)
-        {
-            if (OnClickOverrideHandler != null)
-                OnClickOverrideHandler.Invoke(obj);
-            else
-                InspectorManager.Inspect(obj);
-        }
-
-        public void OnCellExpandToggled(CachedTransform cache)
-        {
-            var instanceID = cache.InstanceID;
-            if (expandedInstanceIDs.Contains(instanceID))
-                expandedInstanceIDs.Remove(instanceID);
-            else
-                expandedInstanceIDs.Add(instanceID);
-
-            RefreshData(true);
-        }
-
-        public void OnCellEnableToggled(CachedTransform cache)
-        {
-            cache.Value.gameObject.SetActive(!cache.Value.gameObject.activeSelf);
-
-            RefreshData(true);
-        }
-
         public void Init()
         {
             ScrollPool.Initialize(this);
@@ -129,7 +97,7 @@ public void JumpAndExpandToTransform(Transform transform)
             }
 
             // Refresh cached transforms (no UI rebuild yet)
-            RefreshData(false);
+            RefreshData(false, false, false);
 
             int transformID = transform.GetInstanceID();
 
@@ -167,57 +135,82 @@ public void Rebuild()
             autoExpandedIDs.Clear();
             expandedInstanceIDs.Clear();
 
-            RefreshData(true, true);
+            RefreshData(true, true, true);
         }
 
-        public void RefreshData(bool andReload = false, bool jumpToTop = false)
+        public void RefreshData(bool andRefreshUI, bool jumpToTop, bool stopExistingCoroutine)
         {
+            if (refreshCoroutine != null)
+            {
+                if (stopExistingCoroutine)
+                {
+                    RuntimeHelper.StopCoroutine(refreshCoroutine);
+                    refreshCoroutine = null;
+                }
+                else
+                    return;
+            }
+
             visited.Clear();
             displayIndex = 0;
             needRefresh = false;
+            traversedThisFrame.Reset();
+            traversedThisFrame.Start();
 
-            var rootObjects = GetRootEntriesMethod.Invoke();
+            IEnumerable<GameObject> rootObjects = GetRootEntriesMethod.Invoke();
 
-            //int displayIndex = 0;
-            foreach (var obj in rootObjects)
-                if (obj) Traverse(obj.transform);
+            refreshCoroutine = RuntimeHelper.StartCoroutine(RefreshCoroutine(rootObjects, andRefreshUI, jumpToTop));
+        }
+
+        private IEnumerator RefreshCoroutine(IEnumerable<GameObject> rootObjects, bool andRefreshUI, bool jumpToTop)
+        {
+            var thisCoro = refreshCoroutine;
+            foreach (var gameObj in rootObjects)
+            {
+                if (gameObj)
+                {
+                    var enumerator = Traverse(gameObj.transform);
+                    while (enumerator.MoveNext())
+                        yield return enumerator.Current;
+                }
+            }
 
             // Prune displayed transforms that we didnt visit in that traverse
             for (int i = cachedTransforms.Count - 1; i >= 0; i--)
             {
-                var obj = (CachedTransform)cachedTransforms[i];
-                if (!visited.Contains(obj.InstanceID))
+                var cached = (CachedTransform)cachedTransforms[i];
+                if (!visited.Contains(cached.InstanceID))
                 {
-                    cachedTransforms.Remove(obj.InstanceID);
+                    cachedTransforms.RemoveAt(i);
                     needRefresh = true;
                 }
             }
 
-            if (!needRefresh)
-                return;
+            if (andRefreshUI && needRefresh)
+                ScrollPool.Refresh(true, jumpToTop);
 
-            //displayedObjects.Clear();
+            prevDisplayIndex = displayIndex;
+        }
 
-            if (andReload)
+        private IEnumerator Traverse(Transform transform, CachedTransform parent = null, int depth = 0)
+        {
+            // Let's only tank 2ms of each frame (60->53fps)
+            if (traversedThisFrame.ElapsedMilliseconds > 2)
             {
-                if (!jumpToTop)
-                    ScrollPool.Refresh(true);
-                else
-                    ScrollPool.Refresh(true, true);
+                yield return null;
+                traversedThisFrame.Reset();
+                traversedThisFrame.Start();
             }
-        }
 
-        private void Traverse(Transform transform, CachedTransform parent = null, int depth = 0)
-        {
             int instanceID = transform.GetInstanceID();
 
             if (visited.Contains(instanceID))
-                return;
+                yield break;
 
             if (Filtering)
             {
                 if (!FilterHierarchy(transform))
-                    return;
+                    yield break;
 
                 visited.Add(instanceID);
 
@@ -231,8 +224,21 @@ private void Traverse(Transform transform, CachedTransform parent = null, int de
             if (cachedTransforms.Contains(instanceID))
             {
                 cached = (CachedTransform)cachedTransforms[(object)instanceID];
+                int prevSiblingIdx = cached.SiblingIndex;
                 if (cached.Update(transform, depth))
+                {
                     needRefresh = true;
+
+                    // If the sibling index changed, we need to shuffle it in our cached transforms list.
+                    if (prevSiblingIdx != cached.SiblingIndex)
+                    {
+                        cachedTransforms.Remove(instanceID);
+                        if (cachedTransforms.Count <= displayIndex)
+                            cachedTransforms.Add(instanceID, cached);
+                        else
+                            cachedTransforms.Insert(displayIndex, instanceID, cached);
+                    }
+                }
             }
             else
             {
@@ -249,7 +255,11 @@ private void Traverse(Transform transform, CachedTransform parent = null, int de
             if (IsCellExpanded(instanceID) && cached.Value.childCount > 0)
             {
                 for (int i = 0; i < transform.childCount; i++)
-                    Traverse(transform.GetChild(i), cached, depth + 1);
+                {
+                    var enumerator = Traverse(transform.GetChild(i), cached, depth + 1);
+                    while (enumerator.MoveNext())
+                        yield return enumerator.Current;
+                }
             }
         }
 
@@ -276,13 +286,44 @@ public void SetCell(TransformCell cell, int index)
                 if (Filtering)
                 {
                     if (cell.cachedTransform.Name.ContainsIgnoreCase(currentFilter))
-                    {
                         cell.NameButton.ButtonText.color = Color.green;
-                    }
                 }
             }
             else
                 cell.Disable();
         }
+
+        public void OnCellBorrowed(TransformCell cell)
+        {
+            cell.OnExpandToggled += OnCellExpandToggled;
+            cell.OnGameObjectClicked += OnGameObjectClicked;
+            cell.OnEnableToggled += OnCellEnableToggled;
+        }
+
+        private void OnGameObjectClicked(GameObject obj)
+        {
+            if (OnClickOverrideHandler != null)
+                OnClickOverrideHandler.Invoke(obj);
+            else
+                InspectorManager.Inspect(obj);
+        }
+
+        public void OnCellExpandToggled(CachedTransform cache)
+        {
+            var instanceID = cache.InstanceID;
+            if (expandedInstanceIDs.Contains(instanceID))
+                expandedInstanceIDs.Remove(instanceID);
+            else
+                expandedInstanceIDs.Add(instanceID);
+
+            RefreshData(true, false, true);
+        }
+
+        public void OnCellEnableToggled(CachedTransform cache)
+        {
+            cache.Value.gameObject.SetActive(!cache.Value.gameObject.activeSelf);
+
+            RefreshData(true, false, true);
+        }
     }
 }