From a286c2ecd9cbd6a891bd7035f6a9bdb85713c1b4 Mon Sep 17 00:00:00 2001
From: Relintai <relintai@protonmail.com>
Date: Sun, 20 Nov 2022 23:15:42 +0100
Subject: [PATCH] Now custom nodes can be added to the SceneTreeDock's create
 root node page.

---
 editor/scene_tree_dock.cpp | 66 +++++++++++++++++++++++++++++++++++---
 editor/scene_tree_dock.h   | 22 ++++++++++---
 2 files changed, 80 insertions(+), 8 deletions(-)

diff --git a/editor/scene_tree_dock.cpp b/editor/scene_tree_dock.cpp
index 9a1b015a0..362ebaf02 100644
--- a/editor/scene_tree_dock.cpp
+++ b/editor/scene_tree_dock.cpp
@@ -1177,11 +1177,12 @@ void SceneTreeDock::_tool_selected(int p_tool, bool p_confirm_override) {
 		case TOOL_CREATE_2D_SCENE:
 		case TOOL_CREATE_3D_SCENE:
 		case TOOL_CREATE_USER_INTERFACE:
+		case TOOL_CREATE_CUSTOM_ROOT_CLASS:
 		case TOOL_CREATE_FAVORITE: {
 			Node *new_node = nullptr;
 
 			if (TOOL_CREATE_FAVORITE == p_tool) {
-				String name = selected_favorite_root.get_slicec(' ', 0);
+				String name = selected_root.get_slicec(' ', 0);
 				if (ScriptServer::is_global_class(name)) {
 					new_node = Object::cast_to<Node>(ClassDB::instance(ScriptServer::get_global_class_native_base(name)));
 					Ref<Script> script = ResourceLoader::load(ScriptServer::get_global_class_path(name), "Script");
@@ -1190,12 +1191,19 @@ void SceneTreeDock::_tool_selected(int p_tool, bool p_confirm_override) {
 						new_node->set_name(name);
 					}
 				} else {
-					new_node = Object::cast_to<Node>(ClassDB::instance(selected_favorite_root));
+					new_node = Object::cast_to<Node>(ClassDB::instance(selected_root));
 				}
 
 				if (!new_node) {
 					new_node = memnew(Node);
-					ERR_PRINT("Creating root from favorite '" + selected_favorite_root + "' failed. Creating 'Node' instead.");
+					ERR_PRINT("Creating root from favorite '" + selected_root + "' failed. Creating 'Node' instead.");
+				}
+			} else if (p_tool == TOOL_CREATE_CUSTOM_ROOT_CLASS) {
+				new_node = Object::cast_to<Node>(ClassDB::instance(selected_root));
+
+				if (!new_node) {
+					new_node = memnew(Node);
+					ERR_PRINT("Creating custom root '" + selected_root + "' failed. Creating 'Node' instead.");
 				}
 			} else {
 				switch (p_tool) {
@@ -3092,6 +3100,19 @@ void SceneTreeDock::_update_create_root_dialog() {
 		return;
 	}
 
+	if (beginner_nodes) {
+		for (int i = 0; i < _custom_scene_root_classes.size(); ++i) {
+			CustomSceneRootClassEntry e = _custom_scene_root_classes[i];
+
+			Button *button = memnew(Button);
+			beginner_nodes->add_child(button);
+			button->set_text(TTR(e.name));
+			button->set_clip_text(true);
+			button->set_icon(EditorNode::get_singleton()->get_class_icon(e.class_name));
+			button->connect("pressed", this, "_custom_root_selected", make_binds(e.class_name));
+		}
+	}
+
 	EditorSettings::get_singleton()->set_setting("_use_favorites_root_selection", toggle->is_pressed());
 	EditorSettings::get_singleton()->save();
 	if (toggle->is_pressed()) {
@@ -3135,10 +3156,15 @@ void SceneTreeDock::_update_create_root_dialog() {
 }
 
 void SceneTreeDock::_favorite_root_selected(const String &p_class) {
-	selected_favorite_root = p_class;
+	selected_root = p_class;
 	_tool_selected(TOOL_CREATE_FAVORITE, false);
 }
 
+void SceneTreeDock::_custom_root_selected(const String &p_class) {
+	selected_root = p_class;
+	_tool_selected(TOOL_CREATE_CUSTOM_ROOT_CLASS, false);
+}
+
 void SceneTreeDock::_clear_clipboard() {
 	for (List<Node *>::Element *E = node_clipboard.front(); E; E = E->next()) {
 		memdelete(E->get());
@@ -3243,6 +3269,7 @@ void SceneTreeDock::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("_local_tree_selected"), &SceneTreeDock::_local_tree_selected);
 	ClassDB::bind_method(D_METHOD("_update_script_button"), &SceneTreeDock::_update_script_button);
 	ClassDB::bind_method(D_METHOD("_favorite_root_selected"), &SceneTreeDock::_favorite_root_selected);
+	ClassDB::bind_method(D_METHOD("_custom_root_selected"), &SceneTreeDock::_custom_root_selected);
 	ClassDB::bind_method(D_METHOD("_update_create_root_dialog"), &SceneTreeDock::_update_create_root_dialog);
 
 	ClassDB::bind_method(D_METHOD("instance"), &SceneTreeDock::instance);
@@ -3254,6 +3281,37 @@ void SceneTreeDock::_bind_methods() {
 	ADD_SIGNAL(MethodInfo("node_created", PropertyInfo(Variant::OBJECT, "node", PROPERTY_HINT_RESOURCE_TYPE, "Node")));
 }
 
+int SceneTreeDock::add_custom_scene_root_class(const String &p_name, const String &p_class_name) {
+	int next_id = 0;
+
+	for (int i = 0; i < _custom_scene_root_classes.size(); ++i) {
+		int cid = _custom_scene_root_classes[i].id;
+
+		if (cid > next_id) {
+			next_id = cid;
+		}
+	}
+
+	++next_id;
+
+	CustomSceneRootClassEntry e;
+	e.id = next_id;
+	e.name = p_name;
+	e.class_name = p_class_name;
+
+	_custom_scene_root_classes.push_back(e);
+
+	return next_id;
+}
+void SceneTreeDock::remove_custom_scene_root_class(const int id) {
+	for (int i = 0; i < _custom_scene_root_classes.size(); ++i) {
+		if (_custom_scene_root_classes[i].id == id) {
+			_custom_scene_root_classes.remove(i);
+			return;
+		}
+	}
+}
+
 SceneTreeDock::SceneTreeDock(EditorNode *p_editor, Node *p_scene_root, EditorSelection *p_editor_selection, EditorData &p_editor_data) {
 	set_name("Scene");
 	editor = p_editor;
diff --git a/editor/scene_tree_dock.h b/editor/scene_tree_dock.h
index 5bb0ec171..bf225b930 100644
--- a/editor/scene_tree_dock.h
+++ b/editor/scene_tree_dock.h
@@ -35,13 +35,13 @@
 #include "core/containers/hash_map.h"
 #include "core/containers/list.h"
 #include "core/containers/map.h"
-#include "core/string/node_path.h"
+#include "core/containers/set.h"
+#include "core/containers/vector.h"
 #include "core/object/object_id.h"
 #include "core/object/reference.h"
 #include "core/object/resource.h"
-#include "core/containers/set.h"
+#include "core/string/node_path.h"
 #include "core/string/ustring.h"
-#include "core/containers/vector.h"
 
 #include "modules/modules_enabled.gen.h" // For regex.
 
@@ -117,6 +117,7 @@ class SceneTreeDock : public VBoxContainer {
 		TOOL_CREATE_2D_SCENE,
 		TOOL_CREATE_3D_SCENE,
 		TOOL_CREATE_USER_INTERFACE,
+		TOOL_CREATE_CUSTOM_ROOT_CLASS,
 		TOOL_CREATE_FAVORITE,
 	};
 
@@ -195,7 +196,7 @@ class SceneTreeDock : public VBoxContainer {
 	bool tree_clicked;
 
 	VBoxContainer *create_root_dialog;
-	String selected_favorite_root;
+	String selected_root;
 
 	void _add_children_to_popup(Object *p_obj, int p_depth);
 
@@ -271,6 +272,7 @@ class SceneTreeDock : public VBoxContainer {
 
 	void _update_create_root_dialog();
 	void _favorite_root_selected(const String &p_class);
+	void _custom_root_selected(const String &p_class);
 
 	void _clear_clipboard();
 	void _create_remap_for_node(Node *p_node, Map<RES, RES> &r_remap);
@@ -279,6 +281,14 @@ class SceneTreeDock : public VBoxContainer {
 	bool _update_node_path(Node *p_root_node, NodePath &r_node_path, Map<Node *, NodePath> *p_renames) const;
 	bool _check_node_path_recursive(Node *p_root_node, Variant &r_variant, Map<Node *, NodePath> *p_renames) const;
 
+	struct CustomSceneRootClassEntry {
+		int id;
+		String name;
+		String class_name;
+	};
+
+	Vector<CustomSceneRootClassEntry> _custom_scene_root_classes;
+
 protected:
 	void _notification(int p_what);
 	static void _bind_methods();
@@ -297,6 +307,7 @@ public:
 	void set_selected(Node *p_node, bool p_emit_selected = false);
 	void fill_path_renames(Node *p_node, Node *p_new_parent, Map<Node *, NodePath> *p_renames);
 	void perform_node_renames(Node *p_base, Map<Node *, NodePath> *p_renames, Map<Ref<Animation>, Set<int>> *r_rem_anims = nullptr);
+
 	SceneTreeEditor *get_tree_editor() {
 		return scene_tree;
 	}
@@ -322,6 +333,9 @@ public:
 		return script_create_dialog;
 	}
 
+	int add_custom_scene_root_class(const String &p_name, const String &p_class_name);
+	void remove_custom_scene_root_class(const int id);
+
 	SceneTreeDock(EditorNode *p_editor, Node *p_scene_root, EditorSelection *p_editor_selection, EditorData &p_editor_data);
 	~SceneTreeDock();
 };