Implement tree search: highlighting

2024-09-25 15:32:03 +00:00 · 2024-09-25 15:32:03 +00:00 · a80b737319
parent 7c4a462a69
commit a80b737319
6 changed files with 455 additions and 57 deletions
--- a/editor/limbo_ai_editor_plugin.cpp
+++ b/editor/limbo_ai_editor_plugin.cpp
@ -457,6 +457,8 @@ void LimboAIEditor::_process_shortcut_input(const Ref<InputEvent> &p_event) {
 			_on_save_pressed();
 		} else if (LW_IS_SHORTCUT("limbo_ai/load_behavior_tree", p_event)) {
 			_popup_file_dialog(load_dialog);
+		} else if (LW_IS_SHORTCUT("limbo_ai/search_tree", p_event)) {
+			task_tree->tree_search_show_and_focus();
 		} else {
 			handled = false;
 		}
@ -800,7 +802,7 @@ void LimboAIEditor::_misc_option_selected(int p_id) {
 			EDIT_SCRIPT(template_path);
 		} break;
 		case MISC_SEARCH_TREE: {
-			UtilityFunctions::print("Search Tree");
+			task_tree->tree_search_show_and_focus();
 		}
 	}
 }
@ -1520,6 +1522,8 @@ LimboAIEditor::LimboAIEditor() {
 	LW_SHORTCUT("limbo_ai/jump_to_owner", TTR("Jump to Owner"), (Key)(LW_KEY_MASK(CMD_OR_CTRL) | LW_KEY(J)));
 	LW_SHORTCUT("limbo_ai/close_tab", TTR("Close Tab"), (Key)(LW_KEY_MASK(CMD_OR_CTRL) | LW_KEY(W)));
 	LW_SHORTCUT("limbo_ai/find_task", TTR("Find Task"), (Key)(LW_KEY_MASK(CMD_OR_CTRL) | LW_KEY(F)));
+	LW_SHORTCUT("limbo_ai/hide_tree_search", TTR("Hide BehaviorTrees Search Panel"), (Key)(LW_KEY(ESCAPE)));
+	LW_SHORTCUT("limbo_ai/search_tree", TTR("Shows the BehaviorTree Search Panel"), (Key)(LW_KEY_MASK(CMD_OR_CTRL) | LW_KEY(F)));

 	set_process_shortcut_input(true);

--- a/editor/task_tree.cpp
+++ b/editor/task_tree.cpp
@ -106,6 +106,7 @@ void TaskTree::_update_item(TreeItem *p_item) {
 	if (!warning_text.is_empty()) {
 		p_item->add_button(0, theme_cache.task_warning_icon, 0, false, warning_text);
 	}
+	tree_search->on_item_edited(p_item); // this is necessary to preserve custom drawing from tree search.
 }

 void TaskTree::_update_tree() {
@ -125,7 +126,7 @@ void TaskTree::_update_tree() {
 	for (const Ref<BTTask> &task : selection) {
 		add_selection(task);
 	}
-	tree_search.apply_search(tree);
+	tree_search->update_search(tree);
 }

 TreeItem *TaskTree::_find_item(const Ref<BTTask> &p_task) const {
@ -532,6 +533,8 @@ void TaskTree::_notification(int p_what) {
 			tree->connect("multi_selected", callable_mp(this, &TaskTree::_on_item_selected).unbind(3), CONNECT_DEFERRED);
 			tree->connect("item_activated", callable_mp(this, &TaskTree::_on_item_activated));
 			tree->connect("item_collapsed", callable_mp(this, &TaskTree::_on_item_collapsed));
+			tree_search->search_panel->connect("text_changed", callable_mp(this, &TaskTree::_update_tree).unbind(1));
+			tree_search->search_panel->connect("visibility_changed", callable_mp(this, &TaskTree::_update_tree));
 		} break;
 		case NOTIFICATION_THEME_CHANGED: {
 			_do_update_theme_item_cache();
@ -564,6 +567,10 @@ void TaskTree::_bind_methods() {
 			PropertyInfo(Variant::INT, "type")));
 }

+void TaskTree::tree_search_show_and_focus() {
+	tree_search->search_panel->show_and_focus();
+}
+
 TaskTree::TaskTree() {
 	editable = true;
 	updating_tree = false;
@ -587,7 +594,9 @@ TaskTree::TaskTree() {
 	tree->set_select_mode(Tree::SelectMode::SELECT_MULTI);

 	tree->set_drag_forwarding(callable_mp(this, &TaskTree::_get_drag_data_fw), callable_mp(this, &TaskTree::_can_drop_data_fw), callable_mp(this, &TaskTree::_drop_data_fw));
-	vbox_container->add_child(tree_search.search_panel);
+	
+	tree_search = memnew(TreeSearch);
+	vbox_container->add_child(tree_search->search_panel);
 }

 TaskTree::~TaskTree() {
--- a/editor/task_tree.h
+++ b/editor/task_tree.h
@ -46,7 +46,7 @@ private:
 	bool editable;
 	bool updating_tree;
 	HashMap<RECT_CACHE_KEY, Rect2> probability_rect_cache;
-	TreeSearch tree_search;
+	TreeSearch * tree_search;

 	struct ThemeCache {
 		Ref<Font> comment_font;
@ -101,6 +101,7 @@ public:
 	Ref<BTTask> get_selected() const;
 	Vector<Ref<BTTask>> get_selected_tasks() const;
 	void clear_selection();
+	void tree_search_show_and_focus();

 	Rect2 get_selected_probability_rect() const;
 	double get_selected_probability_weight() const;
@ -109,6 +110,7 @@ public:

 	virtual bool editor_can_reload_from_file() { return false; }
 	
+
 	TaskTree();
 	~TaskTree();
 };
--- a/editor/tree_search.cpp
+++ b/editor/tree_search.cpp
@ -9,16 +9,35 @@
 * =============================================================================
 */

+#ifdef TOOLS_ENABLED
+
 #include "tree_search.h"

-#include "../util/limbo_compat.h" // for edge scale
+#include "../bt/behavior_tree.h"
+#include "../util/limbo_compat.h" // for edscale
 #include "../util/limbo_string_names.h"
 #include "../util/limbo_utility.h"
-#include <godot_cpp/classes/h_separator.hpp>
-#include <godot_cpp/classes/editor_interface.hpp> // for edge scale
-#include <godot_cpp/templates/vector.hpp>

-#ifdef TOOLS_ENABLED
+#ifdef LIMBOAI_MODULE
+#include "core/math/math_funcs.h"
+#include "editor/editor_interface.h"
+#include "editor/themes/editor_scale.h"
+#include "scene/resources/font.h"
+#include "scene/gui/separator.h"
+#include "scene/resources/style_box_flat.h"
+#include "scene/main/viewport.h"
+#endif // LIMBOAI_MODULE
+
+#ifdef LIMBOAI_GDEXTENSION
+#include <godot_cpp/core/math.hpp>
+#include <godot_cpp/classes/editor_interface.hpp> // for edge scale
+#include <godot_cpp/classes/font.hpp>
+#include <godot_cpp/classes/h_separator.hpp>
+#include <godot_cpp/classes/style_box_flat.hpp>
+#include <godot_cpp/classes/viewport.hpp>
+#endif // LIMBOAI_GDEXTENSION
+
+#define UPPER_BOUND (1 << 15) // for substring search.

 /* ------- TreeSearchPanel ------- */

@ -34,7 +53,9 @@ void TreeSearchPanel::_initialize_controls() {
 	label_highlight->set_text(TTR("Highlight"));
 	label_filter->set_text(TTR("Filter"));

-	close_button->set_button_icon(get_theme_icon(LW_NAME(Close), LW_NAME(EditorIcons)));
+	BUTTON_SET_ICON(close_button, get_theme_icon(LW_NAME(Close), LW_NAME(EditorIcons)));
+	close_button->set_theme_type_variation("FlatButton");
+
 	// positioning and sizing
 	this->set_anchors_and_offsets_preset(LayoutPreset::PRESET_BOTTOM_WIDE);
 	this->set_v_size_flags(SIZE_SHRINK_CENTER); // Do not expand vertically
@ -42,26 +63,59 @@ void TreeSearchPanel::_initialize_controls() {
 	// hack add separator to the left so line edit doesn't clip.
 	line_edit_search->set_h_size_flags(SIZE_EXPAND_FILL);

-	add_spacer(0.25); // otherwise the lineedits expand margin touches the left border.
+	_add_spacer(0.25); // otherwise the lineedits expand margin touches the left border.
 	add_child(line_edit_search);
-	add_spacer();
+	_add_spacer();
 	add_child(label_highlight);
 	add_child(check_button_filter_highlight);
 	add_child(label_filter);
 	add_child(close_button);
-	add_spacer(0.25);
+	_add_spacer(0.25);
 }

-void TreeSearchPanel::add_spacer(float width_multiplier) {
+void TreeSearchPanel::_initialize_close_callbacks() {
+	Callable calleable_set_invisible = Callable(this, "set_visible").bind(false); // don't need a custom bind.
+	close_button->connect("pressed", calleable_set_invisible);
+	close_button->set_shortcut(LW_GET_SHORTCUT("limbo_ai/hide_tree_search"));
+}
+
+void TreeSearchPanel::_add_spacer(float p_width_multiplier) {
 	Control *spacer = memnew(Control);
-	spacer->set_custom_minimum_size(Vector2(8.0 * EDSCALE * width_multiplier, 0.0));
+	spacer->set_custom_minimum_size(Vector2(8.0 * EDSCALE * p_width_multiplier, 0.0));
 	add_child(spacer);
 }

-/* !TreeSearchPanel */
+void TreeSearchPanel::_emit_text_changed(const String &p_text) {
+	this->emit_signal("text_changed", p_text);
+}
+
+void TreeSearchPanel::_emit_text_submitted(const String &p_text) {
+	this->emit_signal("text_submitted");
+}
+
+void TreeSearchPanel::_notification(int p_what) {
+	switch (p_what) {
+		case NOTIFICATION_READY: {
+			_initialize_controls();
+			line_edit_search->connect("text_changed", callable_mp(this, &TreeSearchPanel::_emit_text_changed));
+			_initialize_close_callbacks();
+			line_edit_search->connect("text_submitted", callable_mp(this, &TreeSearchPanel::_emit_text_submitted));
+			break;
+		}
+	}
+}
+
+void TreeSearchPanel::_bind_methods() {
+	ADD_SIGNAL(MethodInfo("text_changed"));
+	ADD_SIGNAL(MethodInfo("text_submitted"));
+}

 TreeSearchPanel::TreeSearchPanel() {
-	_initialize_controls();
+	this->set_visible(false);
+}
+
+bool TreeSearchPanel::has_focus() {
+	return false;
 }

 TreeSearchMode TreeSearchPanel::get_search_mode() {
@ -77,33 +131,324 @@ String TreeSearchPanel::get_text() {
 	return line_edit_search->get_text();
 }

+void TreeSearchPanel::show_and_focus() {
+	this->set_visible(true);
+	line_edit_search->grab_focus();
+}
+
+/* !TreeSearchPanel */
+
 /* ------- TreeSearch ------- */

-void TreeSearch::filter_tree(TreeItem *tree_item, String search_mask) {
-	PRINT_LINE("filter tree not yet implemented!", search_mask);
+void TreeSearch::_filter_tree(TreeItem *p_tree_item, const String &p_search_mask) {
+	for (int i = 0; i < ordered_tree_items.size(); i++){
+		
+	}
+	PRINT_LINE("filter tree not yet implemented!", p_search_mask);
 }

-void TreeSearch::highlight_tree(TreeItem *tree_item, String search_mask) {
-	PRINT_LINE("highlight tree not yet implemented! ", search_mask);
-	// queue/iterative instead of recursive approach. dsf
-	Vector<TreeItem> queue;
-	Vector<TreeItem> hits;
+void TreeSearch::_highlight_tree(const String &p_search_mask) {
+	callable_cache.clear();
+	for (int i = 0; i < ordered_tree_items.size(); i++) {
+		TreeItem *entry = ordered_tree_items[i];
+
+		int num_m = number_matches.has(entry) ? number_matches.get(entry) : 0;
+		if (num_m == 0){
+			continue;;
+		}
+
+		// make sure to also call any draw method already defined.
+		Callable parent_draw_method;
+		if (entry->get_cell_mode(0) == TreeItem::CELL_MODE_CUSTOM) {
+			parent_draw_method = entry->get_custom_draw_callback(0);
+		}
+
+		Callable draw_callback = callable_mp(this, &TreeSearch::_draw_highlight_item).bind(p_search_mask, parent_draw_method);
+	
+		// -- this is necessary because of the modularity of this implementation
+		// cache render properties of entry
+		String cached_text = entry->get_text(0);
+		Ref<Texture2D> cached_icon = entry->get_icon(0);
+		int cached_max_width = entry->get_icon_max_width(0);
+		callable_cache[entry] = draw_callback;
+
+		// this removes render properties in entry
+		entry->set_custom_draw_callback(0, draw_callback);
+		entry->set_cell_mode(0, TreeItem::CELL_MODE_CUSTOM);
+
+		// restore render properties
+		entry->set_text(0, cached_text);
+		entry->set_icon(0, cached_icon);
+		entry->set_icon_max_width(0, cached_max_width);
+	}
 }

-// Call this as a post-processing step for the already constructed tree.
-void TreeSearch::apply_search(Tree *tree) {
-	if (!search_panel || !search_panel->is_visible()){
+// custom draw callback for highlighting (bind 2)
+void TreeSearch::_draw_highlight_item(TreeItem *p_tree_item, Rect2 p_rect, String p_search_mask, Callable p_parent_draw_method) {
+	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);
+	
+	// first part: outline
+	if (matching_entries.has(p_tree_item)) {
+
+		// font info
+		Ref<Font> font = p_tree_item->get_custom_font(0);
+		if (font.is_null()) {
+			font = p_tree_item->get_tree()->get_theme_font(LW_NAME(font));
+		}
+		ERR_FAIL_NULL(font);
+		double font_size = p_tree_item->get_custom_font_size(0);
+		if (font_size == -1) {
+			font_size = p_tree_item->get_tree()->get_theme_font_size(LW_NAME(font));
+		}
+
+		// substring size
+		String string_full = p_tree_item->get_text(0);
+		StringSearchIndices substring_idx = _substring_bounds(string_full, p_search_mask);
+		
+		String substring_match = string_full.substr(substring_idx.lower, substring_idx.upper - substring_idx.lower);
+		Vector2 substring_match_size = font->get_string_size(substring_match, HORIZONTAL_ALIGNMENT_LEFT, -1.f, font_size);
+
+		String substring_before = string_full.substr(0, substring_idx.lower);
+		Vector2 substring_before_size = font->get_string_size(substring_before, HORIZONTAL_ALIGNMENT_LEFT, -1.f, font_size);
+
+		// stylebox
+		Ref<StyleBox> stylebox = p_tree_item->get_tree()->get_theme_stylebox("Focus");
+		ERR_FAIL_NULL(stylebox);
+
+		// extract separation
+		float h_sep = p_tree_item->get_tree()->get_theme_constant("h_separation");
+
+		// compose draw rect
+		const Vector2 PADDING = Vector2(4., 2.);
+		Rect2 draw_rect = p_rect;
+
+		Vector2 rect_offset = Vector2(substring_before_size.x, 0);
+		rect_offset.x += p_tree_item->get_icon_max_width(0) * EDSCALE;
+		rect_offset.x += (h_sep + 4.) * EDSCALE; // TODO: Find better way to determine texts x-offset
+		rect_offset.y = (p_rect.size.y - substring_match_size.y) / 2; // center box vertically
+
+		draw_rect.position += rect_offset - PADDING / 2;
+		draw_rect.size = substring_match_size + PADDING;
+
+		// draw
+		stylebox->draw(p_tree_item->get_tree()->get_canvas_item(), draw_rect);
+	}
+
+	// second part: draw number (TODO: maybe use columns)
+	int num_mat = number_matches.has(p_tree_item) ? number_matches.get(p_tree_item) : 0;
+	if (num_mat > 0) {
+		float h_sep = p_tree_item->get_tree()->get_theme_constant("h_separation");
+		Ref<Font> font = tree_reference->get_theme_font("font");
+		float font_size = tree_reference->get_theme_font_size("font") * 0.75;
+
+		String num_string = String::num_int64(num_mat);
+		Vector2 string_size = font->get_string_size(num_string, HORIZONTAL_ALIGNMENT_CENTER, -1, font_size);
+		Vector2 text_pos = p_rect.position;
+
+		text_pos.x += p_rect.size.x - string_size.x - h_sep;
+		text_pos.y += font->get_descent(font_size) + p_rect.size.y / 2.; // center vertically
+
+		font->draw_string(tree_reference->get_canvas_item(), text_pos, num_string, HORIZONTAL_ALIGNMENT_CENTER, -1, font_size);
+	}
+}
+
+void TreeSearch::_update_matching_entries(const String &search_mask) {
+	Vector<TreeItem *> accum;
+	matching_entries = _find_matching_entries(tree_reference->get_root(), search_mask, accum);
+}
+
+/* this linearizes the tree into [ordered_tree_items] like so:
+ - i1
+	- i2
+	- i3
+ - i4 ---> [i1,i2,i3,i4]
+*/
+void TreeSearch::_update_ordered_tree_items(TreeItem *p_tree_item) {
+	if (!p_tree_item)
+		return;
+	if (p_tree_item == p_tree_item->get_tree()->get_root()) {
+		ordered_tree_items.clear();
+	}
+	// Add the current item to the list
+	ordered_tree_items.push_back(p_tree_item);
+
+	// Recursively collect items from the first child
+	TreeItem *child = p_tree_item->get_first_child();
+	while (child) {
+		_update_ordered_tree_items(child);
+		child = child->get_next();
+	}
+}
+
+void TreeSearch::_update_number_matches() {
+	number_matches.clear();
+	for (int i = 0; i < matching_entries.size(); i++) {
+		TreeItem *item = matching_entries[i];
+		while (item) {
+			int old_num_value = number_matches.has(item) ? number_matches.get(item) : 0;
+			number_matches[item] = old_num_value + 1;
+			item = item->get_parent();
+		}
+	}
+}
+
+Vector<TreeItem *> TreeSearch::_find_matching_entries(TreeItem *p_tree_item, const String &p_search_mask, Vector<TreeItem *> &p_accum) {
+	if (!p_tree_item)
+		return p_accum;
+	StringSearchIndices item_search_indices = _substring_bounds(p_tree_item->get_text(0), p_search_mask);
+	if (item_search_indices.hit())
+		p_accum.insert(p_accum.bsearch(p_tree_item, true), p_tree_item);
+
+	for (int i = 0; i < p_tree_item->get_child_count(); i++) {
+		TreeItem *child = p_tree_item->get_child(i);
+		_find_matching_entries(child, p_search_mask, p_accum);
+	}
+	return p_accum;
+}
+
+// Returns the lower and upper bounds of a substring. Does fuzzy search: Simply looks if words exist in right ordering.
+// Also ignores case if p_search_mask is lowercase. Example:
+// p_searcheable = "TimeLimit 2 sec", p_search_mask = limit 2 sec -> [4,14]. With p_search_mask = "LimiT 2 SEC" or "Limit sec 2" -> [-1,-1]
+TreeSearch::StringSearchIndices TreeSearch::_substring_bounds(const String &p_searchable, const String &p_search_mask) const {
+	StringSearchIndices result;
+	result.lower = UPPER_BOUND;
+	result.upper = 0;
+
+	if (p_search_mask.is_empty()) {
+		return result; // Early return if search_mask is empty.
+	}
+
+	// Determine if the search should be case-insensitive.
+	bool is_case_insensitive = (p_search_mask == p_search_mask.to_lower());
+	String searchable_processed = is_case_insensitive ? p_searchable.to_lower() : p_searchable;
+	PackedStringArray words = p_search_mask.split(" ");
+	int word_position = 0;
+	for (const String &word : words) {
+		if (word.is_empty()) {
+			continue; // Skip empty words.
+		}
+
+		String word_processed = is_case_insensitive ? word.to_lower() : word;
+
+		// Find the position of the next word in the searchable string.
+		word_position = searchable_processed.find(word_processed, word_position);
+
+		if (word_position < 0) {
+			// If any word is not found, return an empty StringSearchIndices.
+			return StringSearchIndices();
+		}
+
+		// Update lower and upper bounds.
+		result.lower = MIN(result.lower, word_position);
+		result.upper = MAX(result.upper, static_cast<int>(word_position + word.length()));
+	}
+
+	return result;
+}
+
+void TreeSearch::_select_item(TreeItem *item) {
+	if (!item)
+		return;
+	tree_reference->set_selected(item, 0);
+	tree_reference->scroll_to_item(item);
+	item = item->get_parent();
+	while (item) {
+		item->set_collapsed(false);
+		item = item->get_parent();
+	}
+}
+
+void TreeSearch::_select_first_match() {
+	if (matching_entries.size() == 0) {
+		return;
+	}
+	for (int i = 0; i < ordered_tree_items.size(); i++) {
+		TreeItem *item = ordered_tree_items[i];
+		int match_idx = matching_entries.bsearch(item, true);
+		if (match_idx < 0 || match_idx >= matching_entries.size() || matching_entries[match_idx] != item) {
+			continue;
+		}
+		String debug_string = "[";
+		_select_item(item);
+		return;
+	}
+}
+
+void TreeSearch::_select_next_match() {
+	if (matching_entries.size() == 0) {
+		return;
+	}
+	TreeItem *selected = tree_reference->get_selected(); // we care about a single item here.
+	if (!selected) {
+		_select_first_match();
 		return;
 	}

-	TreeItem * tree_root = tree->get_root();
-	String search_mask = search_panel->get_text();
-	TreeSearchMode search_mode = search_panel->get_search_mode();
-	if (search_mode == TreeSearchMode::HIGHLIGHT){
-		highlight_tree(tree_root, search_mask);
+	int selected_idx = -1;
+	for (int i = 0; i < ordered_tree_items.size(); i++) {
+		if (ordered_tree_items[i] == selected) {
+			selected_idx = i;
+			break;
 		}
-	if (search_mode == TreeSearchMode::FILTER){
-		filter_tree(tree_root, search_mask);
+	}
+
+	// find the best fitting entry.
+	for (int i = 0; i < ordered_tree_items.size(); i++) {
+		TreeItem *item = ordered_tree_items[i];
+		int match_idx = matching_entries.bsearch(item, true);
+		if (match_idx < 0 || match_idx >= matching_entries.size() || selected_idx >= i || matching_entries[match_idx] != item) {
+			continue;
+		}
+
+		_select_item(item);
+		return;
+	}
+	_select_first_match(); // wrap around.
+}
+
+void TreeSearch::on_item_edited(TreeItem * item) {
+	if (item->get_cell_mode(0) != TreeItem::CELL_MODE_CUSTOM){
+		return;
+	}
+
+	if (!callable_cache.has(item) || item->get_custom_draw_callback(0) == callable_cache.get(item)){
+		return;
+	}
+
+	item->set_custom_draw_callback(0,callable_cache.get(item));
+}
+
+// Call this as a post-processing step for the already constructed tree.
+void TreeSearch::update_search(Tree *p_tree) {
+	ERR_FAIL_COND(!search_panel || !p_tree);
+
+	// ignore if panel not visible or no search string is given.
+	if (!search_panel->is_visible() || search_panel->get_text().length() == 0) {
+		return;
+	}
+
+	tree_reference = p_tree;
+
+	String search_mask = search_panel->get_text();
+	TreeItem *tree_root = p_tree->get_root();
+	TreeSearchMode search_mode = search_panel->get_search_mode();
+
+	_update_ordered_tree_items(p_tree->get_root());
+	_update_matching_entries(search_mask);
+	_update_number_matches();
+
+
+	if (search_mode == TreeSearchMode::HIGHLIGHT) {
+		_highlight_tree(search_mask);
+		if (!search_panel->is_connected("text_submitted", callable_mp(this, &TreeSearch::_select_next_match))) {
+			search_panel->connect("text_submitted", callable_mp(this, &TreeSearch::_select_next_match));
+		}
+	}
+	if (search_mode == TreeSearchMode::FILTER) {
+		_filter_tree(tree_root, search_mask);
 	}
 }

@ -116,4 +461,4 @@ TreeSearch::~TreeSearch() {

 /* !TreeSearch */

-#endif // ! TOOLS_ENABLED
+#endif // TOOLS_ENABLED
--- a/editor/tree_search.h
+++ b/editor/tree_search.h
@ -18,16 +18,19 @@

 #ifdef LIMBOAI_GDEXTENSION
 #include <godot_cpp/classes/check_button.hpp>
-#include <godot_cpp/classes/control.hpp>
 #include <godot_cpp/classes/h_box_container.hpp>
+#include <godot_cpp/classes/label.hpp>
 #include <godot_cpp/classes/line_edit.hpp>
 #include <godot_cpp/classes/tree.hpp>
-#include <godot_cpp/variant/signal.hpp>
-#include <godot_cpp/classes/label.hpp>
+
 #endif // LIMBOAI_GDEXTENSION

 #ifdef LIMBOAI_MODULE
-// TODO: Add includes for godot module variant.
+#include "scene/gui/check_button.h"
+#include "scene/gui/box_container.h"
+#include "scene/gui/label.h"
+#include "scene/gui/line_edit.h"
+#include "scene/gui/tree.h"
 #endif // LIMBOAI_MODULE

 using namespace godot;
@ -38,42 +41,76 @@ enum TreeSearchMode {
 };

 class TreeSearchPanel : public HBoxContainer {
-GDCLASS(TreeSearchPanel, HBoxContainer)
+	GDCLASS(TreeSearchPanel, HBoxContainer)
 private:
 	Button *toggle_button_filter_highlight;
 	Button *close_button;
-	Label * label_highlight;
-	Label * label_filter;
+	Label *label_highlight;
+	Label *label_filter;
 	LineEdit *line_edit_search;
 	CheckButton *check_button_filter_highlight;
-	
 	void _initialize_controls();
-	void add_spacer(float width_multiplier = 1.f);
-	
+	void _initialize_close_callbacks();
+	void _add_spacer(float width_multiplier = 1.f);

+	void _on_draw_highlight(TreeItem *item, Rect2 rect);
+	void _emit_text_changed(const String &text);
+	void _emit_text_submitted(const String &p_text);
+	void _notification(int p_what);

 protected:
-	static void _bind_methods(){}; // we don't need anything exposed.
+	static void _bind_methods();
+	void _process_shortcut_input(const Ref<InputEvent> &p_event);

 public:
 	TreeSearchMode get_search_mode();
 	String get_text();
-	
+	void show_and_focus();
 	TreeSearchPanel();
+	bool has_focus();
 };

-class TreeSearch {
+class TreeSearch : public RefCounted {
+	GDCLASS(TreeSearch, RefCounted)
 private:
-	void filter_tree(TreeItem *tree_item, String search_mask);
-	void highlight_tree(TreeItem *tree_item, String search_mask);
+	struct StringSearchIndices {
+		// initialize to opposite bounds.
+		int lower = -1;
+		int upper = -1;

-	void highlight_item(TreeItem * tree_item, String search_mask);
-	Vector<TreeItem *> find_matching_entries(TreeItem * tree_item, Vector<TreeItem * > buffer = Vector<TreeItem *>());
+		bool hit() {
+			return 0 <= lower && lower < upper;
+		}
+	};
+
+	Tree * tree_reference;
+
+	Vector<TreeItem *> ordered_tree_items;
+	Vector<TreeItem *> matching_entries;
+	HashMap<TreeItem *, int> number_matches;
+	HashMap<TreeItem *, Callable> callable_cache;
+	
+	void _filter_tree(TreeItem *tree_item, const String &search_mask);
+	void _highlight_tree(const String &p_search_mask);
+	void _draw_highlight_item(TreeItem *tree_item, Rect2 rect, String search_mask, Callable parent_draw_method);
+	void _update_matching_entries(const String &search_mask);
+	void _update_ordered_tree_items(TreeItem *tree_item);
+	void _update_number_matches();
+	Vector<TreeItem *> _find_matching_entries(TreeItem *tree_item, const String &search_mask, Vector<TreeItem *> &buffer);
+	StringSearchIndices _substring_bounds(const String &searchable, const String &search_mask) const;
+
+	void _select_item(TreeItem * item);
+	void _select_first_match();
+	void _select_next_match();
+
+
+protected:
+	static void _bind_methods(){}

 public:
 	// we will add everything from TaskTree.h
-	void apply_search(Tree *tree);
-	void set_search_evaluation_method(Callable method);
+	void update_search(Tree *tree);
+	void on_item_edited(TreeItem *item);
 	TreeSearchPanel *search_panel;

 	TreeSearch();
--- a/register_types.cpp
+++ b/register_types.cpp
@ -269,6 +269,7 @@ void initialize_limboai_module(ModuleInitializationLevel p_level) {
 		GDREGISTER_CLASS(LimboAIEditor);
 		GDREGISTER_CLASS(LimboAIEditorPlugin);
 		GDREGISTER_INTERNAL_CLASS(TreeSearchPanel);
+		GDREGISTER_INTERNAL_CLASS(TreeSearch);
 #endif // LIMBOAI_GDEXTENSION

 		EditorPlugins::add_by_type<LimboAIEditorPlugin>();