Refactor node processing in visual shader member dialog

2 files changed, 93 insertions, 120 deletions

diff --git a/editor/plugins/visual_shader_editor_plugin.cpp b/editor/plugins/visual_shader_editor_plugin.cpp
index 826af88571..e228f25c62 100644
--- a/editor/plugins/visual_shader_editor_plugin.cpp
+++ b/editor/plugins/visual_shader_editor_plugin.cpp
@@ -116,7 +116,7 @@ void VisualShaderEditor::clear_custom_types() {
 	}
 }
 
-void VisualShaderEditor::add_custom_type(const String &p_name, const Ref<Script> &p_script, const String &p_description, int p_return_icon_type, const String &p_category, const String &p_subcategory, bool p_highend) {
+void VisualShaderEditor::add_custom_type(const String &p_name, const Ref<Script> &p_script, const String &p_description, int p_return_icon_type, const String &p_category, bool p_highend) {
 
 	ERR_FAIL_COND(!p_name.is_valid_identifier());
 	ERR_FAIL_COND(!p_script.is_valid());
@@ -134,15 +134,15 @@ void VisualShaderEditor::add_custom_type(const String &p_name, const Ref<Script>
 	ao.return_type = p_return_icon_type;
 	ao.description = p_description;
 	ao.category = p_category;
-	ao.sub_category = p_subcategory;
 	ao.highend = p_highend;
 	ao.is_custom = true;
 
 	bool begin = false;
+	String root = p_category.split("/")[0];
 
 	for (int i = 0; i < add_options.size(); i++) {
 		if (add_options[i].is_custom) {
-			if (add_options[i].category == p_category) {
+			if (add_options[i].category == root) {
 				if (!begin) {
 					begin = true;
 				}
@@ -239,9 +239,6 @@ void VisualShaderEditor::update_custom_nodes() {
 			if (ref->has_method("_get_category")) {
 				category = (String)ref->call("_get_category");
 			}
-			if (category == "") {
-				category = "Custom";
-			}
 
 			String subcategory = "";
 			if (ref->has_method("_get_subcategory")) {
@@ -258,18 +255,19 @@ void VisualShaderEditor::update_custom_nodes() {
 			dict["script"] = script;
 			dict["description"] = description;
 			dict["return_icon_type"] = return_icon_type;
+
+			category = category.rstrip("/");
+			category = category.lstrip("/");
+			category = "Addons/" + category;
+			if (subcategory != "") {
+				category += "/" + subcategory;
+			}
+
 			dict["category"] = category;
-			dict["subcategory"] = subcategory;
 			dict["highend"] = highend;
 
 			String key;
-			key = category;
-			key += "/";
-			if (subcategory != "") {
-				key += subcategory;
-				key += "/";
-			}
-			key += name;
+			key = category + "/" + name;
 
 			added[key] = dict;
 		}
@@ -284,7 +282,7 @@ void VisualShaderEditor::update_custom_nodes() {
 
 		const Dictionary &value = (Dictionary)added[key];
 
-		add_custom_type(value["name"], value["script"], value["description"], value["return_icon_type"], value["category"], value["subcategory"], value["highend"]);
+		add_custom_type(value["name"], value["script"], value["description"], value["return_icon_type"], value["category"], value["highend"]);
 	}
 
 	_update_options_menu();
@@ -299,22 +297,12 @@ void VisualShaderEditor::_update_options_menu() {
 	node_desc->set_text("");
 	members_dialog->get_ok()->set_disabled(true);
 
-	String prev_category;
-	String prev_sub_category;
-
 	members->clear();
 	TreeItem *root = members->create_item();
-	TreeItem *category = NULL;
-	TreeItem *sub_category = NULL;
 
 	String filter = node_filter->get_text().strip_edges();
 	bool use_filter = !filter.empty();
 
-	Vector<String> categories;
-	Vector<String> sub_categories;
-
-	int item_count = 0;
-	int item_count2 = 0;
 	bool is_first_item = true;
 
 	Color unsupported_color = get_color("error_color", "Editor");
@@ -322,111 +310,92 @@ void VisualShaderEditor::_update_options_menu() {
 
 	static bool low_driver = ProjectSettings::get_singleton()->get("rendering/quality/driver/driver_name") == "GLES2";
 
+	Map<String, TreeItem *> folders;
+
 	int current_func = -1;
 
 	if (!visual_shader.is_null()) {
 		current_func = visual_shader->get_mode();
 	}
 
-	for (int i = 0; i < add_options.size() + 1; i++) {
+	Vector<AddOption> custom_options;
+	Vector<AddOption> embedded_options;
 
-		if (i == add_options.size()) {
-			if (sub_category != NULL && item_count2 == 0) {
-				memdelete(sub_category);
-				--item_count;
+	for (int i = 0; i < add_options.size(); i++) {
+		if (!use_filter || add_options[i].name.findn(filter) != -1) {
+			if ((add_options[i].func != current_func && add_options[i].func != -1) || !_is_available(add_options[i].mode)) {
+				continue;
 			}
-			if (category != NULL && item_count == 0) {
-				memdelete(category);
+			const_cast<AddOption &>(add_options[i]).temp_idx = i; // save valid id
+			if (add_options[i].is_custom) {
+				custom_options.push_back(add_options[i]);
+			} else {
+				embedded_options.push_back(add_options[i]);
 			}
-			break;
 		}
+	}
+	Vector<AddOption> options;
+	SortArray<AddOption, _OptionComparator> sorter;
+	sorter.sort(custom_options.ptrw(), custom_options.size());
 
-		if (!use_filter || add_options[i].name.findn(filter) != -1) {
-
-			if ((add_options[i].func != current_func && add_options[i].func != -1) || !_is_available(add_options[i].mode))
-				continue;
+	options.append_array(custom_options);
+	options.append_array(embedded_options);
 
-			if (prev_category != add_options[i].category) {
-				if (category != NULL && item_count == 0) {
-					memdelete(category);
-				}
+	for (int i = 0; i < options.size(); i++) {
+		String path = options[i].category;
+		Vector<String> subfolders = path.split("/");
+		TreeItem *category = NULL;
 
-				item_count = 0;
-				prev_sub_category = "";
-				category = members->create_item(root);
-				category->set_text(0, add_options[i].category);
-				category->set_selectable(0, false);
-				if (!use_filter)
-					category->set_collapsed(true);
-			}
-
-			if (add_options[i].sub_category != "") {
-				if (prev_sub_category != add_options[i].sub_category) {
-					if (category != NULL) {
-						if (sub_category != NULL && item_count2 == 0) {
-							memdelete(sub_category);
-							--item_count;
-						}
-						++item_count;
-						item_count2 = 0;
-						sub_category = members->create_item(category);
-						sub_category->set_text(0, add_options[i].sub_category);
-						sub_category->set_selectable(0, false);
-						if (!use_filter)
-							sub_category->set_collapsed(true);
-					}
-				}
-			} else {
-				sub_category = NULL;
-			}
-
-			TreeItem *p_category = NULL;
-
-			if (sub_category != NULL) {
-				p_category = sub_category;
-				++item_count2;
-			} else if (category != NULL) {
-				p_category = category;
-				++item_count;
-			}
-
-			if (p_category != NULL) {
-				TreeItem *item = members->create_item(p_category);
-				if (add_options[i].highend && low_driver)
-					item->set_custom_color(0, unsupported_color);
-				else if (add_options[i].highend)
-					item->set_custom_color(0, supported_color);
-				item->set_text(0, add_options[i].name);
-				if (is_first_item && use_filter) {
-					item->select(0);
-					node_desc->set_text(_get_description(i));
-					is_first_item = false;
-				}
-				switch (add_options[i].return_type) {
-					case VisualShaderNode::PORT_TYPE_SCALAR:
-						item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("float", "EditorIcons"));
-						break;
-					case VisualShaderNode::PORT_TYPE_VECTOR:
-						item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("Vector3", "EditorIcons"));
-						break;
-					case VisualShaderNode::PORT_TYPE_BOOLEAN:
-						item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("bool", "EditorIcons"));
-						break;
-					case VisualShaderNode::PORT_TYPE_TRANSFORM:
-						item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("Transform", "EditorIcons"));
-						break;
-					case VisualShaderNode::PORT_TYPE_SAMPLER:
-						item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("ImageTexture", "EditorIcons"));
-						break;
-					default:
-						break;
+		if (!folders.has(path)) {
+			category = root;
+			String path_temp = "";
+			for (int j = 0; j < subfolders.size(); j++) {
+				path_temp += subfolders[j];
+				if (!folders.has(path_temp)) {
+					category = members->create_item(category);
+					category->set_selectable(0, false);
+					category->set_collapsed(!use_filter);
+					category->set_text(0, subfolders[j]);
+					folders.insert(path_temp, category);
+				} else {
+					category = folders[path_temp];
 				}
-				item->set_meta("id", i);
 			}
-
-			prev_sub_category = add_options[i].sub_category;
-			prev_category = add_options[i].category;
+		} else {
+			category = folders[path];
+		}
+
+		TreeItem *item = members->create_item(category);
+		if (options[i].highend && low_driver)
+			item->set_custom_color(0, unsupported_color);
+		else if (options[i].highend)
+			item->set_custom_color(0, supported_color);
+		item->set_text(0, options[i].name);
+		if (is_first_item && use_filter) {
+			item->select(0);
+			node_desc->set_text(options[i].description);
+			is_first_item = false;
+		}
+		switch (options[i].return_type) {
+			case VisualShaderNode::PORT_TYPE_SCALAR:
+				item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("float", "EditorIcons"));
+				break;
+			case VisualShaderNode::PORT_TYPE_VECTOR:
+				item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("Vector3", "EditorIcons"));
+				break;
+			case VisualShaderNode::PORT_TYPE_BOOLEAN:
+				item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("bool", "EditorIcons"));
+				break;
+			case VisualShaderNode::PORT_TYPE_TRANSFORM:
+				item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("Transform", "EditorIcons"));
+				break;
+			case VisualShaderNode::PORT_TYPE_SAMPLER:
+				item->set_icon(0, EditorNode::get_singleton()->get_gui_base()->get_icon("ImageTexture", "EditorIcons"));
+				break;
+			default:
+				break;
 		}
+		item->set_meta("id", options[i].temp_idx);
 	}
 }
 
diff --git a/editor/plugins/visual_shader_editor_plugin.h b/editor/plugins/visual_shader_editor_plugin.h
index 150cf16167..8919690ada 100644
--- a/editor/plugins/visual_shader_editor_plugin.h
+++ b/editor/plugins/visual_shader_editor_plugin.h
@@ -104,7 +104,6 @@ class VisualShaderEditor : public VBoxContainer {
 	struct AddOption {
 		String name;
 		String category;
-		String sub_category;
 		String type;
 		String description;
 		int sub_func;
@@ -116,12 +115,12 @@ class VisualShaderEditor : public VBoxContainer {
 		float value;
 		bool highend;
 		bool is_custom;
+		int temp_idx;
 
 		AddOption(const String &p_name = String(), const String &p_category = String(), const String &p_sub_category = String(), const String &p_type = String(), const String &p_description = String(), int p_sub_func = -1, int p_return_type = -1, int p_mode = -1, int p_func = -1, float p_value = -1, bool p_highend = false) {
 			name = p_name;
 			type = p_type;
-			category = p_category;
-			sub_category = p_sub_category;
+			category = p_category + "/" + p_sub_category;
 			description = p_description;
 			sub_func = p_sub_func;
 			return_type = p_return_type;
@@ -135,8 +134,7 @@ class VisualShaderEditor : public VBoxContainer {
 		AddOption(const String &p_name, const String &p_category, const String &p_sub_category, const String &p_type, const String &p_description, const String &p_sub_func, int p_return_type = -1, int p_mode = -1, int p_func = -1, float p_value = -1, bool p_highend = false) {
 			name = p_name;
 			type = p_type;
-			category = p_category;
-			sub_category = p_sub_category;
+			category = p_category + "/" + p_sub_category;
 			description = p_description;
 			sub_func = 0;
 			sub_func_str = p_sub_func;
@@ -148,6 +146,12 @@ class VisualShaderEditor : public VBoxContainer {
 			is_custom = false;
 		}
 	};
+	struct _OptionComparator {
+
+		_FORCE_INLINE_ bool operator()(const AddOption &a, const AddOption &b) const {
+			return a.category.count("/") > b.category.count("/") || (a.category + "/" + a.name).naturalnocasecmp_to(b.category + "/" + b.name) < 0;
+		}
+	};
 
 	Vector<AddOption> add_options;
 	int texture_node_option_idx;
@@ -265,7 +269,7 @@ public:
 	static VisualShaderEditor *get_singleton() { return singleton; }
 
 	void clear_custom_types();
-	void add_custom_type(const String &p_name, const Ref<Script> &p_script, const String &p_description, int p_return_icon_type, const String &p_category, const String &p_subcategory, bool p_highend);
+	void add_custom_type(const String &p_name, const Ref<Script> &p_script, const String &p_description, int p_return_icon_type, const String &p_category, bool p_highend);
 
 	virtual Size2 get_minimum_size() const;
 	void edit(VisualShader *p_visual_shader);