Merge pull request #19936 from Xrayez/collapse-button

Expand/collapse nodes recursively in scene tree dock

2 files changed, 88 insertions, 2 deletions

diff --git a/editor/scene_tree_dock.cpp b/editor/scene_tree_dock.cpp
index 8130744b4e..442de08ffa 100644
--- a/editor/scene_tree_dock.cpp
+++ b/editor/scene_tree_dock.cpp
@@ -89,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)) {
@@ -370,6 +372,23 @@ void SceneTreeDock::_tool_selected(int p_tool, bool p_confirm_override) {
 			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: {
 
 			if (!profile_allow_editing) {
@@ -1000,6 +1019,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) {
@@ -1032,6 +1062,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);
@@ -1625,6 +1656,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();
@@ -2244,8 +2321,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();
 	}
 
@@ -2523,6 +2602,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);
@@ -2562,6 +2642,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"));
@@ -2579,7 +2660,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 4ede97dcaa..3729e27ce6 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 {