From 540bb585a8903057475ac173c2b6cb8506f58672 Mon Sep 17 00:00:00 2001
From: Alexander Montag <mail@alexander-montag.de>
Date: Tue, 1 Oct 2024 17:04:46 +0200
Subject: [PATCH] Refactor tree filtering logic for correctness and efficiency

---
 editor/tree_search.cpp | 84 ++++++++++++++++++++++++------------------
 1 file changed, 48 insertions(+), 36 deletions(-)

diff --git a/editor/tree_search.cpp b/editor/tree_search.cpp
index 785cd10..2e9b46d 100644
--- a/editor/tree_search.cpp
+++ b/editor/tree_search.cpp
@@ -25,6 +25,7 @@
 #include "scene/main/viewport.h"
 #include "scene/resources/font.h"
 #include "scene/resources/style_box_flat.h"
+#include "core/templates/hash_set.h"
 #endif // LIMBOAI_MODULE
 
 #ifdef LIMBOAI_GDEXTENSION
@@ -34,6 +35,7 @@
 #include <godot_cpp/classes/style_box_flat.hpp>
 #include <godot_cpp/classes/viewport.hpp>
 #include <godot_cpp/core/math.hpp>
+#include <godot_cpp/templates/hash_set.hpp>
 
 #endif // LIMBOAI_GDEXTENSION
 
@@ -55,51 +57,60 @@ void TreeSearch::_clean_callable_cache() {
 	callable_cache = new_callable_cache;
 }
 
+// Makes all tree items invisible that don't match the following criteria:
+// 1. is parent of matching tree_item
+// 2. is matching tree_item
+// 3. is any descendent of matching tree item
 void TreeSearch::_filter_tree(const String &p_search_mask) {
+	ERR_FAIL_COND(!tree_reference);
+
+	// We don't filter if no search_mask is given.
 	if (matching_entries.is_empty()) {
 		return;
 	}
 
-	Vector<bool> item_visibilities;
-	item_visibilities.resize(ordered_tree_items.size());
-	item_visibilities.fill(false);
+	HashSet<TreeItem *> visible_items;
+	visible_items.reserve(ordered_tree_items.size());
 
-	// Make all entries visible that have any matching descendents. O(n)
-	for (int i = 0; i < ordered_tree_items.size(); i++) {
-		TreeItem *entry = ordered_tree_items[i];
-		if (number_matches.has(entry) && number_matches[entry] > 0) {
-			item_visibilities.set(i, true);
+	Vector<TreeItem *> items = { tree_reference->get_root() };
+
+	for (int idx = 0; idx < items.size(); idx++) {
+		TreeItem *cur_item = items[idx];
+		if (_vector_has_bsearch(matching_entries, cur_item)) {
+			visible_items.insert(cur_item);
+			TreeItem *parent_of_match = cur_item->get_parent();
+			
+			// 1. mark all parents visible
+			while (parent_of_match) {
+				visible_items.insert(parent_of_match);
+				parent_of_match = parent_of_match->get_parent();
+			}
+
+			// for 2. and 3. we collect all children of match.
+			Vector<TreeItem *> matching_entries_descendents = { cur_item };
+
+			// 2. and 3. mark all descendents visible
+			for (int i = 0; i < matching_entries_descendents.size(); i++) {
+				TreeItem *cur_descendent_item = matching_entries_descendents[i];
+				visible_items.insert(cur_descendent_item);
+				
+				// 3. collect all children of descendent
+				for (int i = 0; i < cur_descendent_item->get_child_count(); i++) {
+					matching_entries_descendents.push_back(cur_descendent_item->get_child(i));
+				}
+			}
+		
+		} else { // not a matching entry: process children
+			for (int i = 0; i < cur_item->get_child_count(); i++) {
+				items.push_back(cur_item->get_child(i));
+			}
 		}
 	}
 
-	// Make all descendents of matching entries visible. O(n) * log(|matching_entries|)
 	for (int i = 0; i < ordered_tree_items.size(); i++) {
-		TreeItem *entry = ordered_tree_items[i];
-		if (_vector_has_bsearch(matching_entries, entry)) {
-			// the [next_entry] at the same depth or depth above.
-			TreeItem *next_entry = entry->get_next();
-
-			// search above current depth if no [next_entry].
-			while (!next_entry && entry) {
-				entry = entry->get_parent();
-			}
-			if (entry) {
-				next_entry = entry->get_next();
-			}
-
-			// mark visible all successors upto next entry at the same depth or above
-			int j = i;
-			for (; j < ordered_tree_items.size() && ordered_tree_items[j] != next_entry; j++) {
-				item_visibilities.set(j, true);
-			}
-			i = j - 1; // every entry is only processed once.
-		}
-	}
-
-	// Apply visibility. O(n)
-	for (int i = 0; i < ordered_tree_items.size(); i++) {
-		if (!item_visibilities[i]) {
-			ordered_tree_items[i]->set_visible(false);
+		TreeItem *cur_item = ordered_tree_items[i];
+		if (!visible_items.has(cur_item)) {
+			cur_item->set_visible(false);
 		}
 	}
 }
@@ -170,6 +181,7 @@ void TreeSearch::_draw_highlight_item(TreeItem *p_tree_item, const Rect2 p_rect,
 	if (!p_tree_item) {
 		return;
 	}
+
 	// call any parent draw methods such as for probability FIRST.
 	p_parent_draw_method.call(p_tree_item, p_rect);
 
@@ -584,4 +596,4 @@ void TreeSearchPanel::focus_editor() {
 
 /* !TreeSearchPanel */
 
-#endif // TOOLS_ENABLED
\ No newline at end of file
+#endif // TOOLS_ENABLED