Add `_get_func_code/_is_available` virtual functions to custom nodes

6 files changed, 165 insertions, 46 deletions

diff --git a/doc/classes/VisualShaderNodeCustom.xml b/doc/classes/VisualShaderNodeCustom.xml
index 7b992abe24..0a962a4aa4 100644
--- a/doc/classes/VisualShaderNodeCustom.xml
+++ b/doc/classes/VisualShaderNodeCustom.xml
@@ -44,6 +44,17 @@
 				Defining this method is [b]optional[/b].
 			</description>
 		</method>
+		<method name="_get_func_code" qualifiers="virtual const">
+			<return type="String" />
+			<argument index="0" name="mode" type="int" enum="Shader.Mode" />
+			<argument index="1" name="type" type="int" enum="VisualShader.Type" />
+			<description>
+				Override this method to add a shader code to the beginning of each shader function (once). The shader code should be returned as a string, which can have multiple lines (the [code]"""[/code] multiline string construct can be used for convenience).
+				If there are multiple custom nodes of different types which use this feature the order of each insertion is undefined.
+				You can customize the generated code based on the shader [code]mode[/code] (see [enum Shader.Mode]) and/or [code]type[/code] (see [enum VisualShader.Type]).
+				Defining this method is [b]optional[/b].
+			</description>
+		</method>
 		<method name="_get_global_code" qualifiers="virtual const">
 			<return type="String" />
 			<argument index="0" name="mode" type="int" enum="Shader.Mode" />
@@ -114,11 +125,20 @@
 				Defining this method is [b]optional[/b]. If not overridden, no return icon is shown.
 			</description>
 		</method>
+		<method name="_is_available" qualifiers="virtual const">
+			<return type="bool" />
+			<argument index="0" name="mode" type="int" enum="Shader.Mode" />
+			<argument index="1" name="type" type="int" enum="VisualShader.Type" />
+			<description>
+				Override this method to prevent the node to be visible in the member dialog for the certain [code]mode[/code] (see [enum Shader.Mode]) and/or [code]type[/code] (see [enum VisualShader.Type]).
+				Defining this method is [b]optional[/b]. If not overridden, it's [code]true[/code].
+			</description>
+		</method>
 		<method name="_is_highend" qualifiers="virtual const">
 			<return type="bool" />
 			<description>
 				Override this method to enable high-end mark in the Visual Shader Editor's members dialog.
-				Defining this method is [b]optional[/b]. If not overridden, it's false.
+				Defining this method is [b]optional[/b]. If not overridden, it's [code]false[/code].
 			</description>
 		</method>
 	</methods>
diff --git a/editor/plugins/visual_shader_editor_plugin.cpp b/editor/plugins/visual_shader_editor_plugin.cpp
index 94906acb63..7f30dd91e5 100644
--- a/editor/plugins/visual_shader_editor_plugin.cpp
+++ b/editor/plugins/visual_shader_editor_plugin.cpp
@@ -1082,6 +1082,10 @@ void VisualShaderEditor::edit(VisualShader *p_visual_shader) {
 	}
 }
 
+void VisualShaderEditor::update_nodes() {
+	_update_nodes();
+}
+
 void VisualShaderEditor::add_plugin(const Ref<VisualShaderNodePlugin> &p_plugin) {
 	if (plugins.has(p_plugin)) {
 		return;
@@ -1165,10 +1169,7 @@ bool VisualShaderEditor::_is_available(int p_mode) {
 	return (p_mode == -1 || (p_mode & current_mode) != 0);
 }
 
-void VisualShaderEditor::update_custom_nodes() {
-	if (members_dialog->is_visible()) {
-		return;
-	}
+void VisualShaderEditor::_update_nodes() {
 	clear_custom_types();
 	List<StringName> class_list;
 	ScriptServer::get_global_class_list(&class_list);
@@ -1184,6 +1185,9 @@ void VisualShaderEditor::update_custom_nodes() {
 			Ref<VisualShaderNodeCustom> ref;
 			ref.instantiate();
 			ref->set_script(script);
+			if (!ref->is_available(visual_shader->get_mode(), visual_shader->get_shader_type())) {
+				continue;
+			}
 
 			String name;
 			if (ref->has_method("_get_name")) {
@@ -1240,6 +1244,32 @@ void VisualShaderEditor::update_custom_nodes() {
 		}
 	}
 
+	// Disables not-supported copied items.
+	{
+		for (CopyItem &item : copy_items_buffer) {
+			Ref<VisualShaderNodeCustom> custom = Object::cast_to<VisualShaderNodeCustom>(item.node.ptr());
+
+			if (custom.is_valid()) {
+				if (!custom->is_available(visual_shader->get_mode(), visual_shader->get_shader_type())) {
+					item.disabled = true;
+				} else {
+					item.disabled = false;
+				}
+			} else {
+				for (int i = 0; i < add_options.size(); i++) {
+					if (add_options[i].type == item.node->get_class_name()) {
+						if ((add_options[i].func != visual_shader->get_mode() && add_options[i].func != -1) || !_is_available(add_options[i].mode)) {
+							item.disabled = true;
+						} else {
+							item.disabled = false;
+						}
+						break;
+					}
+				}
+			}
+		}
+	}
+
 	Array keys = added.keys();
 	keys.sort();
 
@@ -3370,15 +3400,23 @@ void VisualShaderEditor::_graph_gui_input(const Ref<InputEvent> &p_event) {
 			selected_float_constant = -1;
 		}
 
-		if (to_change.is_empty() && copy_items_buffer.is_empty()) {
+		bool copy_buffer_empty = true;
+		for (const CopyItem &item : copy_items_buffer) {
+			if (!item.disabled) {
+				copy_buffer_empty = false;
+				break;
+			}
+		}
+
+		if (to_change.is_empty() && copy_buffer_empty) {
 			_show_members_dialog(true);
 		} else {
 			popup_menu->set_item_disabled(NodeMenuOptions::CUT, to_change.is_empty());
 			popup_menu->set_item_disabled(NodeMenuOptions::COPY, to_change.is_empty());
-			popup_menu->set_item_disabled(NodeMenuOptions::PASTE, copy_items_buffer.is_empty());
+			popup_menu->set_item_disabled(NodeMenuOptions::PASTE, copy_buffer_empty);
 			popup_menu->set_item_disabled(NodeMenuOptions::DELETE, to_change.is_empty());
 			popup_menu->set_item_disabled(NodeMenuOptions::DUPLICATE, to_change.is_empty());
-			popup_menu->set_item_disabled(NodeMenuOptions::CLEAR_COPY_BUFFER, copy_items_buffer.is_empty());
+			popup_menu->set_item_disabled(NodeMenuOptions::CLEAR_COPY_BUFFER, copy_buffer_empty);
 
 			int temp = popup_menu->get_item_index(NodeMenuOptions::SEPARATOR2);
 			if (temp != -1) {
@@ -3715,6 +3753,17 @@ void VisualShaderEditor::_dup_paste_nodes(int p_type, List<CopyItem> &r_items, c
 	if (p_duplicate) {
 		undo_redo->create_action(TTR("Duplicate VisualShader Node(s)"));
 	} else {
+		bool copy_buffer_empty = true;
+		for (const CopyItem &item : copy_items_buffer) {
+			if (!item.disabled) {
+				copy_buffer_empty = false;
+				break;
+			}
+		}
+		if (copy_buffer_empty) {
+			return;
+		}
+
 		undo_redo->create_action(TTR("Paste VisualShader Node(s)"));
 	}
 
@@ -3727,16 +3776,7 @@ void VisualShaderEditor::_dup_paste_nodes(int p_type, List<CopyItem> &r_items, c
 	Set<int> added_set;
 
 	for (CopyItem &item : r_items) {
-		bool unsupported = false;
-		for (int i = 0; i < add_options.size(); i++) {
-			if (add_options[i].type == item.node->get_class_name()) {
-				if (!_is_available(add_options[i].mode)) {
-					unsupported = true;
-				}
-				break;
-			}
-		}
-		if (unsupported) {
+		if (item.disabled) {
 			unsupported_set.insert(item.id);
 			continue;
 		}
@@ -3777,7 +3817,10 @@ void VisualShaderEditor::_dup_paste_nodes(int p_type, List<CopyItem> &r_items, c
 	}
 
 	id_from = base_id;
-	for (int i = 0; i < r_items.size(); i++) {
+	for (const CopyItem &item : r_items) {
+		if (item.disabled) {
+			continue;
+		}
 		undo_redo->add_undo_method(visual_shader.ptr(), "remove_node", type, id_from);
 		undo_redo->add_undo_method(graph_plugin.ptr(), "remove_node", type, id_from);
 		id_from++;
@@ -3878,7 +3921,7 @@ void VisualShaderEditor::_mode_selected(int p_id) {
 	}
 
 	visual_shader->set_shader_type(VisualShader::Type(p_id + offset));
-	_update_options_menu();
+	_update_nodes();
 	_update_graph();
 
 	graph->grab_focus();
@@ -4465,8 +4508,8 @@ void VisualShaderEditor::_visibility_changed() {
 }
 
 void VisualShaderEditor::_bind_methods() {
+	ClassDB::bind_method("_update_nodes", &VisualShaderEditor::_update_nodes);
 	ClassDB::bind_method("_update_graph", &VisualShaderEditor::_update_graph);
-	ClassDB::bind_method("_update_options_menu", &VisualShaderEditor::_update_options_menu);
 	ClassDB::bind_method("_add_node", &VisualShaderEditor::_add_node);
 	ClassDB::bind_method("_node_changed", &VisualShaderEditor::_node_changed);
 	ClassDB::bind_method("_input_select_item", &VisualShaderEditor::_input_select_item);
@@ -5439,7 +5482,7 @@ void VisualShaderEditorPlugin::make_visible(bool p_visible) {
 		//editor->animation_panel_make_visible(true);
 		button->show();
 		EditorNode::get_singleton()->make_bottom_panel_item_visible(visual_shader_editor);
-		visual_shader_editor->update_custom_nodes();
+		visual_shader_editor->update_nodes();
 		visual_shader_editor->set_process_input(true);
 		//visual_shader_editor->set_process(true);
 	} else {
@@ -5919,8 +5962,8 @@ void EditorPropertyShaderMode::_option_selected(int p_which) {
 		}
 	}
 
-	undo_redo->add_do_method(editor, "_update_options_menu");
-	undo_redo->add_undo_method(editor, "_update_options_menu");
+	undo_redo->add_do_method(editor, "_update_nodes");
+	undo_redo->add_undo_method(editor, "_update_nodes");
 
 	undo_redo->add_do_method(editor, "_update_graph");
 	undo_redo->add_undo_method(editor, "_update_graph");
diff --git a/editor/plugins/visual_shader_editor_plugin.h b/editor/plugins/visual_shader_editor_plugin.h
index e26b606397..4087fc583c 100644
--- a/editor/plugins/visual_shader_editor_plugin.h
+++ b/editor/plugins/visual_shader_editor_plugin.h
@@ -262,6 +262,7 @@ class VisualShaderEditor : public VBoxContainer {
 	void _show_add_varying_dialog();
 	void _show_remove_varying_dialog();
 
+	void _update_nodes();
 	void _update_graph();
 
 	struct AddOption {
@@ -394,6 +395,7 @@ class VisualShaderEditor : public VBoxContainer {
 		String group_inputs;
 		String group_outputs;
 		String expression;
+		bool disabled = false;
 	};
 
 	void _dup_copy_nodes(int p_type, List<CopyItem> &r_nodes, List<VisualShader::Connection> &r_connections);
@@ -476,7 +478,7 @@ protected:
 	static void _bind_methods();
 
 public:
-	void update_custom_nodes();
+	void update_nodes();
 	void add_plugin(const Ref<VisualShaderNodePlugin> &p_plugin);
 	void remove_plugin(const Ref<VisualShaderNodePlugin> &p_plugin);
 
diff --git a/scene/resources/visual_shader.cpp b/scene/resources/visual_shader.cpp
index 03b768e869..d3e61dbc84 100644
--- a/scene/resources/visual_shader.cpp
+++ b/scene/resources/visual_shader.cpp
@@ -455,32 +455,69 @@ String VisualShaderNodeCustom::generate_code(Shader::Mode p_mode, VisualShader::
 	for (int i = 0; i < get_output_port_count(); i++) {
 		output_vars.push_back(p_output_vars[i]);
 	}
-	String code = "	{\n";
+
 	String _code;
 	GDVIRTUAL_CALL(_get_code, input_vars, output_vars, p_mode, p_type, _code);
-	bool nend = _code.ends_with("\n");
-	_code = _code.insert(0, "		");
-	_code = _code.replace("\n", "\n		");
-	code += _code;
-	if (!nend) {
-		code += "\n	}";
-	} else {
-		code.remove_at(code.size() - 1);
-		code += "}";
+	if (_is_valid_code(_code)) {
+		String code = "	{\n";
+		bool nend = _code.ends_with("\n");
+		_code = _code.insert(0, "		");
+		_code = _code.replace("\n", "\n		");
+		code += _code;
+		if (!nend) {
+			code += "\n	}";
+		} else {
+			code.remove_at(code.size() - 1);
+			code += "}";
+		}
+		code += "\n";
+		return code;
 	}
-	code += "\n";
-	return code;
+	return String();
 }
 
 String VisualShaderNodeCustom::generate_global_per_node(Shader::Mode p_mode, int p_id) const {
-	String ret;
-	if (GDVIRTUAL_CALL(_get_global_code, p_mode, ret)) {
-		String code = "// " + get_caption() + "\n";
-		code += ret;
-		code += "\n";
-		return code;
+	String _code;
+	if (GDVIRTUAL_CALL(_get_global_code, p_mode, _code)) {
+		if (_is_valid_code(_code)) {
+			String code = "// " + get_caption() + "\n";
+			code += _code;
+			code += "\n";
+			return code;
+		}
 	}
-	return "";
+	return String();
+}
+
+String VisualShaderNodeCustom::generate_global_per_func(Shader::Mode p_mode, VisualShader::Type p_type, int p_id) const {
+	String _code;
+	if (GDVIRTUAL_CALL(_get_func_code, p_mode, p_type, _code)) {
+		if (_is_valid_code(_code)) {
+			bool nend = _code.ends_with("\n");
+			String code = "// " + get_caption() + "\n";
+			code += "	{\n";
+			_code = _code.insert(0, "	");
+			_code = _code.replace("\n", "\n		");
+			code += _code;
+			if (!nend) {
+				code += "\n	}";
+			} else {
+				code.remove_at(code.size() - 1);
+				code += "}";
+			}
+			code += "\n";
+			return code;
+		}
+	}
+	return String();
+}
+
+bool VisualShaderNodeCustom::is_available(Shader::Mode p_mode, VisualShader::Type p_type) const {
+	bool ret;
+	if (GDVIRTUAL_CALL(_is_available, p_mode, p_type, ret)) {
+		return ret;
+	}
+	return true;
 }
 
 void VisualShaderNodeCustom::set_input_port_default_value(int p_port, const Variant &p_value, const Variant &p_prev_value) {
@@ -511,6 +548,13 @@ void VisualShaderNodeCustom::_set_input_port_default_value(int p_port, const Var
 	VisualShaderNode::set_input_port_default_value(p_port, p_value);
 }
 
+bool VisualShaderNodeCustom::_is_valid_code(const String &p_code) const {
+	if (p_code.is_empty() || p_code == "null") {
+		return false;
+	}
+	return true;
+}
+
 bool VisualShaderNodeCustom::_is_initialized() {
 	return is_initialized;
 }
@@ -531,8 +575,10 @@ void VisualShaderNodeCustom::_bind_methods() {
 	GDVIRTUAL_BIND(_get_output_port_type, "port");
 	GDVIRTUAL_BIND(_get_output_port_name, "port");
 	GDVIRTUAL_BIND(_get_code, "input_vars", "output_vars", "mode", "type");
+	GDVIRTUAL_BIND(_get_func_code, "mode", "type");
 	GDVIRTUAL_BIND(_get_global_code, "mode");
 	GDVIRTUAL_BIND(_is_highend);
+	GDVIRTUAL_BIND(_is_available, "mode", "type");
 
 	ClassDB::bind_method(D_METHOD("_set_initialized", "enabled"), &VisualShaderNodeCustom::_set_initialized);
 	ClassDB::bind_method(D_METHOD("_is_initialized"), &VisualShaderNodeCustom::_is_initialized);
diff --git a/scene/resources/visual_shader.h b/scene/resources/visual_shader.h
index 5cad5fc95b..2d4b2852e9 100644
--- a/scene/resources/visual_shader.h
+++ b/scene/resources/visual_shader.h
@@ -383,14 +383,20 @@ protected:
 	GDVIRTUAL1RC(int, _get_output_port_type, int)
 	GDVIRTUAL1RC(String, _get_output_port_name, int)
 	GDVIRTUAL4RC(String, _get_code, TypedArray<String>, TypedArray<String>, Shader::Mode, VisualShader::Type)
+	GDVIRTUAL2RC(String, _get_func_code, Shader::Mode, VisualShader::Type)
 	GDVIRTUAL1RC(String, _get_global_code, Shader::Mode)
 	GDVIRTUAL0RC(bool, _is_highend)
+	GDVIRTUAL2RC(bool, _is_available, Shader::Mode, VisualShader::Type)
+
+	bool _is_valid_code(const String &p_code) const;
 
 protected:
 	void _set_input_port_default_value(int p_port, const Variant &p_value);
 
+	bool is_available(Shader::Mode p_mode, VisualShader::Type p_type) const;
 	virtual String generate_code(Shader::Mode p_mode, VisualShader::Type p_type, int p_id, const String *p_input_vars, const String *p_output_vars, bool p_for_preview = false) const override;
 	virtual String generate_global_per_node(Shader::Mode p_mode, int p_id) const override;
+	virtual String generate_global_per_func(Shader::Mode p_mode, VisualShader::Type p_type, int p_id) const override;
 
 	static void _bind_methods();
 
diff --git a/scene/resources/visual_shader_nodes.cpp b/scene/resources/visual_shader_nodes.cpp
index a6aa6d8c49..7b4d7e66d4 100644
--- a/scene/resources/visual_shader_nodes.cpp
+++ b/scene/resources/visual_shader_nodes.cpp
@@ -5696,7 +5696,7 @@ String VisualShaderNodeTextureUniformTriplanar::get_input_port_name(int p_port)
 String VisualShaderNodeTextureUniformTriplanar::generate_global_per_node(Shader::Mode p_mode, int p_id) const {
 	String code;
 
-	code += "// TRIPLANAR FUNCTION GLOBAL CODE\n";
+	code += "// " + get_caption() + "\n";
 	code += "	vec4 triplanar_texture(sampler2D p_sampler, vec3 p_weights, vec3 p_triplanar_pos) {\n";
 	code += "		vec4 samp = vec4(0.0);\n";
 	code += "		samp += texture(p_sampler, p_triplanar_pos.xy) * p_weights.z;\n";
@@ -5719,11 +5719,13 @@ String VisualShaderNodeTextureUniformTriplanar::generate_global_per_func(Shader:
 	String code;
 
 	if (p_type == VisualShader::TYPE_VERTEX) {
-		code += "	// TRIPLANAR FUNCTION VERTEX CODE\n";
+		code += "// " + get_caption() + "\n";
+		code += "	{\n";
 		code += "		triplanar_power_normal = pow(abs(NORMAL), vec3(triplanar_sharpness));\n";
 		code += "		triplanar_power_normal /= dot(triplanar_power_normal, vec3(1.0));\n";
 		code += "		triplanar_pos = VERTEX * triplanar_scale + triplanar_offset;\n";
 		code += "		triplanar_pos *= vec3(1.0, -1.0, 1.0);\n";
+		code += "	}\n";
 	}
 
 	return code;