diff --git a/editor/tree_search.cpp b/editor/tree_search.cpp
index 2e9b46d..18fae15 100644
--- a/editor/tree_search.cpp
+++ b/editor/tree_search.cpp
@@ -19,13 +19,13 @@
 
 #ifdef LIMBOAI_MODULE
 #include "core/math/math_funcs.h"
+#include "core/templates/hash_set.h"
 #include "editor/editor_interface.h"
 #include "editor/themes/editor_scale.h"
 #include "scene/gui/separator.h"
 #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
@@ -58,60 +58,55 @@ void TreeSearch::_clean_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
+// 1. is a matching tree_item
+// 2. is a parent of a matching tree_item
+// 3. is any descendant of a 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()) {
+	if (number_matches.size() == 0){
 		return;
 	}
+	HashSet<TreeItem *> visible_as_parent;
+	HashSet<TreeItem *> visible_as_descendant;
 
-	HashSet<TreeItem *> visible_items;
-	visible_items.reserve(ordered_tree_items.size());
+	// Mark matching items and all their ancestors (1, 2)
+	for (int i = 0; i < matching_entries.size(); i++) {
+		TreeItem *cur_match = matching_entries[i];
+		TreeItem *current = cur_match;
+		while (current != nullptr && !visible_as_parent.has(current)) {
+			visible_as_parent.insert(current);
+			current = current->get_parent();
+		}
+	}
 
-	Vector<TreeItem *> items = { tree_reference->get_root() };
+	// Mark matching items and all their descendants (1, 3)
+	for (int i = 0; i < matching_entries.size(); i++) {
+		TreeItem *cur_match = matching_entries[i];
+		if (visible_as_descendant.has(cur_match)) {
+			continue;
+		}
+		// Descendents
+		Vector<TreeItem *> descendent_list = { cur_match };
 
-	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();
+		for (int j = 0; j < descendent_list.size(); j++) {
+			TreeItem *descendent = descendent_list[j];
+			if (visible_as_descendant.has(descendent)) {
+				continue; // Skip subtree if already processed
+			}
+			visible_as_descendant.insert(descendent);
 			
-			// 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));
+			// Collect X-level descendents (bit clunky because godot doesn't return get_children as pointers)
+			for (int k = 0; k < descendent->get_child_count(); k++) {
+				TreeItem *child = descendent->get_child(k);
+				descendent_list.push_back(child);
 			}
 		}
 	}
 
+	// Set visibility based on whether items are in either visibility set
 	for (int i = 0; i < ordered_tree_items.size(); i++) {
-		TreeItem *cur_item = ordered_tree_items[i];
-		if (!visible_items.has(cur_item)) {
-			cur_item->set_visible(false);
-		}
+		TreeItem *item = ordered_tree_items[i];
+		bool is_visible = visible_as_parent.has(item) || visible_as_descendant.has(item);
+		item->set_visible(is_visible);
 	}
 }