From 3e320517910d69208648e87c44a667c92dbd1d3d Mon Sep 17 00:00:00 2001
From: Serhii Snitsaruk <limbonaut@pm.me>
Date: Fri, 2 Feb 2024 16:46:26 +0100
Subject: [PATCH] Editor: Remember collapsed items in task tree

---
 bt/tasks/bt_task.cpp |  8 ++++++++
 bt/tasks/bt_task.h   |  4 ++++
 editor/task_tree.cpp | 22 ++++++++++++++++++++++
 editor/task_tree.h   |  2 ++
 4 files changed, 36 insertions(+)

diff --git a/bt/tasks/bt_task.cpp b/bt/tasks/bt_task.cpp
index 76754472..57481299 100644
--- a/bt/tasks/bt_task.cpp
+++ b/bt/tasks/bt_task.cpp
@@ -88,6 +88,14 @@ void BTTask::_set_children(Array p_children) {
 	}
 }
 
+void BTTask::set_display_collapsed(bool p_display_collapsed) {
+	data.display_collapsed = p_display_collapsed;
+}
+
+bool BTTask::is_displayed_collapsed() const {
+	return data.display_collapsed;
+}
+
 String BTTask::get_task_name() {
 	if (data.custom_name.is_empty()) {
 #ifdef LIMBOAI_MODULE
diff --git a/bt/tasks/bt_task.h b/bt/tasks/bt_task.h
index 95c5593e..720e6500 100644
--- a/bt/tasks/bt_task.h
+++ b/bt/tasks/bt_task.h
@@ -80,6 +80,7 @@ class BTTask : public BT {
 		Vector<Ref<BTTask>> children;
 		Status status = FRESH;
 		double elapsed = 0.0;
+		bool display_collapsed = false;
 	} data;
 
 	Array _get_children() const;
@@ -115,6 +116,9 @@ class BTTask : public BT {
 	_FORCE_INLINE_ Node *get_agent() const { return data.agent; }
 	void set_agent(Node *p_agent) { data.agent = p_agent; }
 
+	void set_display_collapsed(bool p_display_collapsed);
+	bool is_displayed_collapsed() const;
+
 	String get_custom_name() const { return data.custom_name; }
 	void set_custom_name(const String &p_name);
 	String get_task_name();
diff --git a/editor/task_tree.cpp b/editor/task_tree.cpp
index b258417c..dd05e226 100644
--- a/editor/task_tree.cpp
+++ b/editor/task_tree.cpp
@@ -81,6 +81,7 @@ void TaskTree::_update_item(TreeItem *p_item) {
 	p_item->set_icon(0, LimboUtility::get_singleton()->get_task_icon(type_arg));
 	p_item->set_icon_max_width(0, 16 * EDSCALE);
 	p_item->set_editable(0, false);
+	p_item->set_collapsed(task->is_displayed_collapsed());
 
 	for (int i = 0; i < p_item->get_button_count(0); i++) {
 		p_item->erase_button(0, i);
@@ -111,7 +112,9 @@ void TaskTree::_update_tree() {
 	}
 
 	if (bt->get_root_task().is_valid()) {
+		updating_tree = true;
 		_create_tree(bt->get_root_task(), nullptr);
+		updating_tree = false;
 	}
 
 	TreeItem *item = _find_item(sel);
@@ -167,6 +170,21 @@ void TaskTree::_on_item_activated() {
 	emit_signal(LW_NAME(task_activated));
 }
 
+void TaskTree::_on_item_collapsed(Object *p_obj) {
+	if (updating_tree) {
+		return;
+	}
+
+	TreeItem *item = Object::cast_to<TreeItem>(p_obj);
+	if (!item) {
+		return;
+	}
+
+	Ref<BTTask> task = item->get_metadata(0);
+	ERR_FAIL_NULL(task);
+	task->set_display_collapsed(item->is_collapsed());
+}
+
 void TaskTree::_on_task_changed() {
 	_update_item(tree->get_selected());
 }
@@ -183,7 +201,9 @@ void TaskTree::load_bt(const Ref<BehaviorTree> &p_behavior_tree) {
 	tree->clear();
 	probability_rect_cache.clear();
 	if (bt->get_root_task().is_valid()) {
+		updating_tree = true;
 		_create_tree(bt->get_root_task(), nullptr);
+		updating_tree = false;
 	}
 }
 
@@ -366,6 +386,7 @@ void TaskTree::_notification(int p_what) {
 			tree->connect("item_mouse_selected", callable_mp(this, &TaskTree::_on_item_mouse_selected));
 			tree->connect("item_selected", callable_mp(this, &TaskTree::_on_item_selected));
 			tree->connect("item_activated", callable_mp(this, &TaskTree::_on_item_activated));
+			tree->connect("item_collapsed", callable_mp(this, &TaskTree::_on_item_collapsed));
 		} break;
 		case NOTIFICATION_THEME_CHANGED: {
 			_do_update_theme_item_cache();
@@ -399,6 +420,7 @@ void TaskTree::_bind_methods() {
 
 TaskTree::TaskTree() {
 	editable = true;
+	updating_tree = false;
 
 	tree = memnew(Tree);
 	add_child(tree);
diff --git a/editor/task_tree.h b/editor/task_tree.h
index 5515fb6b..d12b28c5 100644
--- a/editor/task_tree.h
+++ b/editor/task_tree.h
@@ -40,6 +40,7 @@ class TaskTree : public Control {
 	Ref<BehaviorTree> bt;
 	Ref<BTTask> last_selected;
 	bool editable;
+	bool updating_tree;
 	HashMap<RECT_CACHE_KEY, Rect2> probability_rect_cache;
 
 	struct ThemeCache {
@@ -67,6 +68,7 @@ class TaskTree : public Control {
 
 	void _on_item_selected();
 	void _on_item_activated();
+	void _on_item_collapsed(Object *p_obj);
 	void _on_item_mouse_selected(const Vector2 &p_pos, MouseButton p_button_index);
 	void _on_task_changed();