From 824f68483f60934f5f6dca11e90f187ab99d57ba Mon Sep 17 00:00:00 2001
From: "Andrii Doroshenko (Xrayez)" <xrayez@gmail.com>
Date: Sat, 7 Jul 2018 17:51:18 +0300
Subject: [PATCH] Add ability to expand/collapse nodes recursively in scene
 tree dock

To expand or collapse the node recursively (all children), hold `Shift`
button and click on the node's folding arrow.

The popup menu option `Expand/Collapse All" checks whether any node
is expanded or collapsed first and performs the opposite operation. That
means if any children node is collapsed, it will first expand all nodes
at selected node.

Co-authored-by: Rikhardur Bjarni Einarsson (MunWolf) badulf96@gmail.com
---
 editor/scene_tree_dock.cpp | 86 +++++++++++++++++++++++++++++++++++++-
 editor/scene_tree_dock.h   |  5 +++
 2 files changed, 89 insertions(+), 2 deletions(-)

diff --git a/editor/scene_tree_dock.cpp b/editor/scene_tree_dock.cpp
index 8473230758d..07544e8b238 100644
--- a/editor/scene_tree_dock.cpp
+++ b/editor/scene_tree_dock.cpp
@@ -31,6 +31,7 @@
 #include "scene_tree_dock.h"
 
 #include "core/io/resource_saver.h"
+#include "core/os/input.h"
 #include "core/os/keyboard.h"
 #include "core/project_settings.h"
 
@@ -88,6 +89,8 @@ void SceneTreeDock::_unhandled_key_input(Ref<InputEvent> p_event) {
 		_tool_selected(TOOL_NEW);
 	} else if (ED_IS_SHORTCUT("scene_tree/instance_scene", p_event)) {
 		_tool_selected(TOOL_INSTANCE);
+	} else if (ED_IS_SHORTCUT("scene_tree/expand_collapse_all", p_event)) {
+		_tool_selected(TOOL_EXPAND_COLLAPSE);
 	} else if (ED_IS_SHORTCUT("scene_tree/change_node_type", p_event)) {
 		_tool_selected(TOOL_REPLACE);
 	} else if (ED_IS_SHORTCUT("scene_tree/duplicate", p_event)) {
@@ -341,6 +344,23 @@ void SceneTreeDock::_tool_selected(int p_tool, bool p_confirm_override) {
 			quick_open->popup_dialog("PackedScene", true);
 			quick_open->set_title(TTR("Instance Child Scene"));
 
+		} break;
+		case TOOL_EXPAND_COLLAPSE: {
+
+			if (!scene_tree->get_selected())
+				break;
+
+			Tree *tree = scene_tree->get_scene_tree();
+			TreeItem *selected_item = tree->get_selected();
+
+			if (!selected_item)
+				selected_item = tree->get_root();
+
+			bool collapsed = _is_collapsed_recursive(selected_item);
+			_set_collapsed_recursive(selected_item, !collapsed);
+
+			tree->ensure_cursor_is_visible();
+
 		} break;
 		case TOOL_REPLACE: {
 
@@ -965,6 +985,17 @@ void SceneTreeDock::_tool_selected(int p_tool, bool p_confirm_override) {
 	}
 }
 
+void SceneTreeDock::_node_collapsed(Object *p_obj) {
+
+	TreeItem *ti = Object::cast_to<TreeItem>(p_obj);
+	if (!ti)
+		return;
+
+	if (Input::get_singleton()->is_key_pressed(KEY_SHIFT)) {
+		_set_collapsed_recursive(ti, ti->is_collapsed());
+	}
+}
+
 void SceneTreeDock::_notification(int p_what) {
 
 	switch (p_what) {
@@ -997,6 +1028,7 @@ void SceneTreeDock::_notification(int p_what) {
 			filter->set_clear_button_enabled(true);
 
 			EditorNode::get_singleton()->get_editor_selection()->connect("selection_changed", this, "_selection_changed");
+			scene_tree->get_scene_tree()->connect("item_collapsed", this, "_node_collapsed");
 
 			// create_root_dialog
 			HBoxContainer *top_row = memnew(HBoxContainer);
@@ -1591,6 +1623,52 @@ void SceneTreeDock::_do_reparent(Node *p_new_parent, int p_position_in_parent, V
 	editor_data->get_undo_redo().commit_action();
 }
 
+bool SceneTreeDock::_is_collapsed_recursive(TreeItem *p_item) const {
+
+	bool is_branch_collapsed = false;
+
+	List<TreeItem *> needs_check;
+	needs_check.push_back(p_item);
+
+	while (!needs_check.empty()) {
+
+		TreeItem *item = needs_check.back()->get();
+		needs_check.pop_back();
+
+		TreeItem *child = item->get_children();
+		is_branch_collapsed = item->is_collapsed() && child;
+
+		if (is_branch_collapsed) {
+			break;
+		}
+		while (child) {
+			needs_check.push_back(child);
+			child = child->get_next();
+		}
+	}
+	return is_branch_collapsed;
+}
+
+void SceneTreeDock::_set_collapsed_recursive(TreeItem *p_item, bool p_collapsed) {
+
+	List<TreeItem *> to_collapse;
+	to_collapse.push_back(p_item);
+
+	while (!to_collapse.empty()) {
+
+		TreeItem *item = to_collapse.back()->get();
+		to_collapse.pop_back();
+
+		item->set_collapsed(p_collapsed);
+
+		TreeItem *child = item->get_children();
+		while (child) {
+			to_collapse.push_back(child);
+			child = child->get_next();
+		}
+	}
+}
+
 void SceneTreeDock::_script_created(Ref<Script> p_script) {
 
 	List<Node *> selected = editor_selection->get_selected_node_list();
@@ -2202,8 +2280,10 @@ void SceneTreeDock::_tree_rmb(const Vector2 &p_menu_pos) {
 
 			menu->add_icon_shortcut(get_icon("Add", "EditorIcons"), ED_GET_SHORTCUT("scene_tree/add_child_node"), TOOL_NEW);
 			menu->add_icon_shortcut(get_icon("Instance", "EditorIcons"), ED_GET_SHORTCUT("scene_tree/instance_scene"), TOOL_INSTANCE);
-			menu->add_separator();
 		}
+		menu->add_icon_shortcut(get_icon("Collapse", "EditorIcons"), ED_GET_SHORTCUT("scene_tree/expand_collapse_all"), TOOL_EXPAND_COLLAPSE);
+		menu->add_separator();
+
 		existing_script = selected->get_script();
 	}
 
@@ -2479,6 +2559,7 @@ void SceneTreeDock::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("_node_prerenamed"), &SceneTreeDock::_node_prerenamed);
 	ClassDB::bind_method(D_METHOD("_import_subscene"), &SceneTreeDock::_import_subscene);
 	ClassDB::bind_method(D_METHOD("_selection_changed"), &SceneTreeDock::_selection_changed);
+	ClassDB::bind_method(D_METHOD("_node_collapsed"), &SceneTreeDock::_node_collapsed);
 	ClassDB::bind_method(D_METHOD("_new_scene_from"), &SceneTreeDock::_new_scene_from);
 	ClassDB::bind_method(D_METHOD("_nodes_dragged"), &SceneTreeDock::_nodes_dragged);
 	ClassDB::bind_method(D_METHOD("_files_dropped"), &SceneTreeDock::_files_dropped);
@@ -2517,6 +2598,7 @@ SceneTreeDock::SceneTreeDock(EditorNode *p_editor, Node *p_scene_root, EditorSel
 	ED_SHORTCUT("scene_tree/batch_rename", TTR("Batch Rename"), KEY_MASK_CMD | KEY_F2);
 	ED_SHORTCUT("scene_tree/add_child_node", TTR("Add Child Node"), KEY_MASK_CMD | KEY_A);
 	ED_SHORTCUT("scene_tree/instance_scene", TTR("Instance Child Scene"));
+	ED_SHORTCUT("scene_tree/expand_collapse_all", TTR("Expand/Collapse All"));
 	ED_SHORTCUT("scene_tree/change_node_type", TTR("Change Type"));
 	ED_SHORTCUT("scene_tree/attach_script", TTR("Attach Script"));
 	ED_SHORTCUT("scene_tree/extend_script", TTR("Extend Script"));
@@ -2534,7 +2616,7 @@ SceneTreeDock::SceneTreeDock(EditorNode *p_editor, Node *p_scene_root, EditorSel
 
 	button_add = memnew(ToolButton);
 	button_add->connect("pressed", this, "_tool_selected", make_binds(TOOL_NEW, false));
-	button_add->set_tooltip(TTR("Add/Create a New Node"));
+	button_add->set_tooltip(TTR("Add/Create a New Node."));
 	button_add->set_shortcut(ED_GET_SHORTCUT("scene_tree/add_child_node"));
 	filter_hbc->add_child(button_add);
 
diff --git a/editor/scene_tree_dock.h b/editor/scene_tree_dock.h
index e66525d721f..6473e261f19 100644
--- a/editor/scene_tree_dock.h
+++ b/editor/scene_tree_dock.h
@@ -60,6 +60,7 @@ class SceneTreeDock : public VBoxContainer {
 
 		TOOL_NEW,
 		TOOL_INSTANCE,
+		TOOL_EXPAND_COLLAPSE,
 		TOOL_RENAME,
 		TOOL_BATCH_RENAME,
 		TOOL_REPLACE,
@@ -116,6 +117,7 @@ class SceneTreeDock : public VBoxContainer {
 
 	HBoxContainer *tool_hbc;
 	void _tool_selected(int p_tool, bool p_confirm_override = false);
+	void _node_collapsed(Object *p_obj);
 
 	EditorData *editor_data;
 	EditorSelection *editor_selection;
@@ -152,6 +154,9 @@ class SceneTreeDock : public VBoxContainer {
 	void _node_reparent(NodePath p_path, bool p_keep_global_xform);
 	void _do_reparent(Node *p_new_parent, int p_position_in_parent, Vector<Node *> p_nodes, bool p_keep_global_xform);
 
+	bool _is_collapsed_recursive(TreeItem *p_item) const;
+	void _set_collapsed_recursive(TreeItem *p_item, bool p_collapsed);
+
 	void _set_owners(Node *p_owner, const Array &p_nodes);
 
 	enum ReplaceOwnerMode {