Fixed bug brough by commit 43442d0690223a2c83e1e2380d94203a9ec52ddc

Editor/Scripts/AbstractTreeView.cs

@@ -2,7 +2,8 @@
 // Heap Explorer for Unity. Copyright (c) 2019-2020 Peter Schraut (www.console-dev.de). See LICENSE.md
 // https://github.com/pschraut/UnityHeapExplorer/
 //
-using System.Collections;
+
+using System;
 using System.Collections.Generic;
 using UnityEngine;
 using UnityEditor.IMGUI.Controls;
@@ -31,12 +32,11 @@ public HeapExplorerWindow window
         SearchTextParser.Result m_Search = new SearchTextParser.Result();
         protected string m_EditorPrefsKey;
         int m_FirstVisibleRow;
+        private List<TreeViewItem> m_RowsCache;
         IList<int> m_Expanded = new List<int>(32);
         TreeViewItem m_Tree;
-        TreeViewItem m_ActiveTree;
         string[] m_SearchCache = new string[32];
         System.Text.StringBuilder m_SearchBuilder = new System.Text.StringBuilder();
-        bool m_HasSearch;
 
         public AbstractTreeView(HeapExplorerWindow window, string editorPrefsKey, TreeViewState state)
             : base(state)
@@ -73,7 +73,6 @@ public AbstractTreeView(HeapExplorerWindow window, string editorPrefsKey, TreeVi
         public void SetTree(TreeViewItem tree)
         {
             m_Tree = tree;
-            m_ActiveTree = m_Tree;
             Reload();
         }
 
@@ -84,107 +83,109 @@ protected override bool CanMultiSelect(TreeViewItem item)
 
         protected override TreeViewItem BuildRoot()
         {
-            if (m_ActiveTree != null)
-                return m_ActiveTree;
+            if (m_Tree != null)
+                return m_Tree;
 
             var root = new TreeViewItem { id = 0, depth = -1, displayName = "Root" };
             root.AddChild(new TreeViewItem { id = root.id + 1, depth = -1, displayName = "" });
             return root;
         }
 
-        void OnSortingChanged(MultiColumnHeader multiColumnHeader)
+        protected override IList<TreeViewItem> BuildRows(TreeViewItem root)
         {
-            if (rootItem == null || !rootItem.hasChildren)
-                return;
+            if (m_RowsCache == null)
+                m_RowsCache = new List<TreeViewItem>(128);
+            m_RowsCache.Clear();
 
-            SortItemsRecursive(m_ActiveTree, OnSortItem);
-            Reload();
-        }
+            if (hasSearch)
+            {
+                SearchTree(root, searchString, m_RowsCache);
+                m_RowsCache.Sort(CompareItem);
+            }
+            else
+            {
+                SortAndAddExpandedRows(root, m_RowsCache);
+            }
 
-        protected abstract int OnSortItem(TreeViewItem x, TreeViewItem y);
+            return m_RowsCache;
+        }
 
-        protected void SortItemsRecursive(TreeViewItem parent, System.Comparison<TreeViewItem> comparison)
+        protected virtual void SearchTree(TreeViewItem root, string search, List<TreeViewItem> result)
         {
-            if (parent == null)
-                return;
+            var stack = new Stack<TreeViewItem>();
+            stack.Push(root);
+            while (stack.Count > 0)
+            {
+                TreeViewItem current = stack.Pop();
+                if (current.children != null)
+                {
+                    foreach (var child in current.children)
+                    {
+                        if (child != null)
+                        {
+                            if (DoesItemMatchSearch(child, search))
+                                result.Add(child);
 
-            var sortMe = new List<TreeViewItem>();
-            sortMe.Add(parent);
+                            stack.Push(child);
+                        }
+                    }
+                }
+            }
+        }
 
-            for (var n = 0; n < sortMe.Count; ++n)
+        protected virtual void SortAndAddExpandedRows(TreeViewItem root, IList<TreeViewItem> rows)
+        {
+            if (root.hasChildren)
             {
-                var item = sortMe[n];
-                if (item.hasChildren)
+                root.children.Sort(CompareItem);
+                foreach (TreeViewItem child in root.children)
                 {
-                    item.children.Sort(comparison);
-                    sortMe.AddRange(item.children); // sort items of children too (kind of recursive)
+                    GetAndSortExpandedRowsRecursive(child, rows);
                 }
             }
         }
 
-        public void Search(string search)
+        void GetAndSortExpandedRowsRecursive(TreeViewItem item, IList<TreeViewItem> expandedRows)
         {
-            m_HasSearch = !string.IsNullOrEmpty(search);
-
-            var selectedId = -1;
-            var selection = new List<int>(this.GetSelection());
-            if (selection != null && selection.Count != 0)
-                selectedId = selection[0];
-            selection.Clear();
-
-            m_Search = SearchTextParser.Parse(search);
+            if (item == null)
+                Debug.LogError("Found a TreeViewItem that is null. Invalid use of AddExpandedRows(): This method is only valid to call if you have built the full tree of TreeViewItems.");
 
-            // I don't use the TreeViewControl base.searchString API, because I didn't get
-            // it to work to display filtered content sorted.
-            // Therefore I perform the filter myself as seen below, by building a new tree
-            // containing a flat list of filtered items only.
-            //base.searchString = search;
+            expandedRows.Add(item);
 
-            if (!m_HasSearch)
-            {
-                m_ActiveTree = m_Tree;
-            }
-            else
+            if (item.hasChildren && IsExpanded(item.id))
             {
-                // Create a node that contains all items that match the search
-                m_ActiveTree = new TreeViewItem { id = 0, depth = -1, displayName = "Root" };
-
-                // Get all items in the tree as a flat list
-                var itemsToProcess = new List<TreeViewItem>();
-                itemsToProcess.Add(m_Tree);
-                for (var n = 0; n < itemsToProcess.Count; ++n)
+                item.children.Sort(CompareItem);
+                foreach (TreeViewItem child in item.children)
                 {
-                    var item = itemsToProcess[n];
-                    if (item == null)
-                        continue;
-
-                    if (item.hasChildren)
-                        itemsToProcess.AddRange(item.children); // process all children too (this makes it kind of recursive)
+                    GetAndSortExpandedRowsRecursive(child, expandedRows);
+                }
+            }
+        }
 
-                    if (item.id == 0)
-                        continue; // ignore tree root node itself
+        void OnSortingChanged(MultiColumnHeader multiColumnHeader)
+        {
+            if (rootItem == null || !rootItem.hasChildren)
+                return;
 
-                    if (!DoesItemMatchSearch(item, search))
-                        continue;
+            Reload();
+        }
 
-                    // Keep track of the selected item
-                    if (selectedId == item.id)
-                    {
-                        selection = new List<int>();
-                        selection.Add(selectedId);
-                    }
+        protected abstract int OnSortItem(TreeViewItem x, TreeViewItem y);
 
-                    m_ActiveTree.AddChild(item);
-                }
+        protected int CompareItem(TreeViewItem x, TreeViewItem y)
+        {
+            int result = OnSortItem(x, y);
+            if (result == 0)
+                return x.id.CompareTo(y.id);
+            return result;
+        }
 
-                // If the search didn't find any item, make sure to add at least one
-                // invisible item, otherwise the TreeViewControl displays an error.
-                if (!m_ActiveTree.hasChildren)
-                    m_ActiveTree.AddChild(new TreeViewItem { id = m_ActiveTree.id + 1, depth = -1, displayName = "" });
-            }
+        public void Search(string search)
+        {
+            var selection = new List<int>(this.GetSelection());
 
-            SortItemsRecursive(m_ActiveTree, OnSortItem);
-            Reload();
+            m_Search = SearchTextParser.Parse(search);
+            base.searchString = search;
 
             if (selection != null && selection.Count > 0)
                 this.SetSelection(selection, TreeViewSelectionOptions.FireSelectionChanged | TreeViewSelectionOptions.RevealAndFrame);
@@ -362,7 +363,7 @@ protected override void RowGUI(RowGUIArgs args)
                     rect.width -= extraSpaceBeforeIconAndLabel;
 
                     // Display the tree as a flat list when content is filtered
-                    if (m_HasSearch)
+                    if (hasSearch)
                         rect = TreeViewUtility.IndentByDepth(0, rect);
                     else
                         rect = TreeViewUtility.IndentByDepth(args.item.depth, rect);

Editor/Scripts/CompareSnapshotsView/CompareSnapshotsControl.cs

@@ -116,8 +116,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshotA, PackedMemorySnapsh
             BuildMemoryTree(snapshots, ref uniqueId, root);
             BuildGCHandleTree(snapshots, ref uniqueId, root);
 
-            SortItemsRecursive(root, OnSortItem);
-
             return root;
         }
 

Editor/Scripts/GCHandlesView/GCHandlesControl.cs

@@ -154,8 +154,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshot)
                 }
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             return root;
         }
 

Editor/Scripts/ManagedHeapSectionsView/ManagedHeapSectionsControl.cs

@@ -105,8 +105,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshot, PackedMemorySection
                 root.AddChild(item);
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             return root;
         }
 

Editor/Scripts/ManagedObjectsView/ManagedObjectsControl.cs

@@ -149,8 +149,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshot)
                 }
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             return root;
         }
 

Editor/Scripts/ManagedTypesView/ManagedTypesControl.cs

@@ -121,8 +121,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshot)
                 }
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             return root;
         }
 

Editor/Scripts/NativeObjectsView/NativeObjectsControl.cs

@@ -253,8 +253,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshot, BuildArgs buildArgs
                 }
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             if (!root.hasChildren)
                 root.AddChild(new TreeViewItem { id = 1, depth = -1, displayName = "" });
 
@@ -349,8 +347,6 @@ public TreeViewItem BuildDuplicatesTree(PackedMemorySnapshot snapshot, BuildArgs
                 group.AddChild(item);
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             if (!root.hasChildren)
                 root.AddChild(new TreeViewItem { id = 1, depth = -1, displayName = "" });
 

Editor/Scripts/StaticFieldsView/StaticFieldsControl.cs

@@ -121,8 +121,6 @@ public TreeViewItem BuildTree(PackedMemorySnapshot snapshot)
                 root.AddChild(group);
             }
 
-            SortItemsRecursive(root, OnSortItem);
-
             return root;
         }