Implement drop-down list properties to the custom visual shader nodes

5 files changed, 230 insertions, 2 deletions

diff --git a/doc/classes/VisualShaderNodeCustom.xml b/doc/classes/VisualShaderNodeCustom.xml
index 8a90d5dd0f..5db0dfb327 100644
--- a/doc/classes/VisualShaderNodeCustom.xml
+++ b/doc/classes/VisualShaderNodeCustom.xml
@@ -80,6 +80,14 @@
 				Defining this method is [b]required[/b]. If not overridden, the node has no input ports.
 			</description>
 		</method>
+		<method name="_get_input_port_default_value" qualifiers="virtual const">
+			<return type="Variant" />
+			<param index="0" name="port" type="int" />
+			<description>
+				Override this method to define the default value for the specified input port. Prefer use this over [method VisualShaderNode.set_input_port_default_value].
+				Defining this method is [b]required[/b]. If not overridden, the node has no default values for their input ports.
+			</description>
+		</method>
 		<method name="_get_input_port_name" qualifiers="virtual const">
 			<return type="String" />
 			<param index="0" name="port" type="int" />
@@ -126,6 +134,37 @@
 				Defining this method is [b]optional[/b], but recommended. If not overridden, output ports will return the [constant VisualShaderNode.PORT_TYPE_SCALAR] type.
 			</description>
 		</method>
+		<method name="_get_property_count" qualifiers="virtual const">
+			<return type="int" />
+			<description>
+				Override this method to define the number of the properties.
+				Defining this method is [b]optional[/b].
+			</description>
+		</method>
+		<method name="_get_property_default_index" qualifiers="virtual const">
+			<return type="int" />
+			<param index="0" name="index" type="int" />
+			<description>
+				Override this method to define the default index of the property of the associated custom node.
+				Defining this method is [b]optional[/b].
+			</description>
+		</method>
+		<method name="_get_property_name" qualifiers="virtual const">
+			<return type="String" />
+			<param index="0" name="index" type="int" />
+			<description>
+				Override this method to define the names of the property of the associated custom node.
+				Defining this method is [b]optional[/b].
+			</description>
+		</method>
+		<method name="_get_property_options" qualifiers="virtual const">
+			<return type="PackedStringArray" />
+			<param index="0" name="index" type="int" />
+			<description>
+				Override this method to define the options inside the drop-down list property of the associated custom node.
+				Defining this method is [b]optional[/b].
+			</description>
+		</method>
 		<method name="_get_return_icon_type" qualifiers="virtual const">
 			<return type="int" enum="VisualShaderNode.PortType" />
 			<description>
@@ -149,5 +188,12 @@
 				Defining this method is [b]optional[/b]. If not overridden, it's [code]false[/code].
 			</description>
 		</method>
+		<method name="get_option_index" qualifiers="const">
+			<return type="int" />
+			<param index="0" name="option" type="int" />
+			<description>
+				Returns the selected index of the drop-down list option within a graph. You may use this function to define the specific behavior in the [method _get_code] or [method _get_global_code].
+			</description>
+		</method>
 	</methods>
 </class>
diff --git a/editor/plugins/visual_shader_editor_plugin.cpp b/editor/plugins/visual_shader_editor_plugin.cpp
index 25cbbbf6de..8b1053b79b 100644
--- a/editor/plugins/visual_shader_editor_plugin.cpp
+++ b/editor/plugins/visual_shader_editor_plugin.cpp
@@ -551,6 +551,47 @@ void VisualShaderGraphPlugin::add_node(VisualShader::Type p_type, int p_id, bool
 		}
 	}
 
+	if (custom_node.is_valid()) {
+		bool first = true;
+		VBoxContainer *vbox = nullptr;
+
+		for (int i = 0; i < custom_node->dp_props.size(); i++) {
+			const VisualShaderNodeCustom::DropDownListProperty &dp = custom_node->dp_props[i];
+
+			if (first) {
+				first = false;
+				vbox = memnew(VBoxContainer);
+				node->add_child(vbox);
+				port_offset++;
+			}
+
+			HBoxContainer *hbox = memnew(HBoxContainer);
+			vbox->add_child(hbox);
+			hbox->set_h_size_flags(Control::SIZE_EXPAND_FILL);
+
+			String prop_name = dp.name.strip_edges();
+			if (!prop_name.is_empty()) {
+				Label *label = memnew(Label);
+				label->set_text(prop_name + ":");
+				hbox->add_child(label);
+			}
+
+			OptionButton *op = memnew(OptionButton);
+			hbox->add_child(op);
+			op->set_h_size_flags(Control::SIZE_EXPAND_FILL);
+			op->connect("item_selected", callable_mp(editor, &VisualShaderEditor::_set_custom_node_option).bind(p_id, i), CONNECT_DEFERRED);
+
+			for (const String &s : dp.options) {
+				op->add_item(s);
+			}
+			if (custom_node->dp_selected_cache.has(i)) {
+				op->select(custom_node->dp_selected_cache[i]);
+			} else {
+				op->select(0);
+			}
+		}
+	}
+
 	Ref<VisualShaderNodeCurveTexture> curve = vsnode;
 	Ref<VisualShaderNodeCurveXYZTexture> curve_xyz = vsnode;
 
@@ -2704,6 +2745,22 @@ void VisualShaderEditor::_edit_port_default_input(Object *p_button, int p_node,
 	editing_port = p_port;
 }
 
+void VisualShaderEditor::_set_custom_node_option(int p_index, int p_node, int p_op) {
+	VisualShader::Type type = get_current_shader_type();
+	Ref<VisualShaderNodeCustom> node = visual_shader->get_node(type, p_node);
+	if (node.is_null()) {
+		return;
+	}
+
+	EditorUndoRedoManager *undo_redo = EditorUndoRedoManager::get_singleton();
+	undo_redo->create_action(TTR("Set Custom Node Option"));
+	undo_redo->add_do_method(node.ptr(), "_set_option_index", p_op, p_index);
+	undo_redo->add_undo_method(node.ptr(), "_set_option_index", p_op, node->get_option_index(p_op));
+	undo_redo->add_do_method(graph_plugin.ptr(), "update_node", type, p_node);
+	undo_redo->add_undo_method(graph_plugin.ptr(), "update_node", type, p_node);
+	undo_redo->commit_action();
+}
+
 void VisualShaderEditor::_setup_node(VisualShaderNode *p_node, const Vector<Variant> &p_ops) {
 	// INPUT
 	{
@@ -3084,7 +3141,9 @@ void VisualShaderEditor::_add_node(int p_idx, const Vector<Variant> &p_ops, Stri
 		}
 		VisualShaderNodeCustom *custom_node = Object::cast_to<VisualShaderNodeCustom>(vsn);
 		ERR_FAIL_NULL(custom_node);
-		custom_node->update_ports();
+		custom_node->update_property_default_values();
+		custom_node->update_input_port_default_values();
+		custom_node->update_properties();
 	}
 
 	bool is_texture2d = (Object::cast_to<VisualShaderNodeTexture>(vsnode.ptr()) != nullptr);
diff --git a/editor/plugins/visual_shader_editor_plugin.h b/editor/plugins/visual_shader_editor_plugin.h
index e0a0f3a096..8629e64467 100644
--- a/editor/plugins/visual_shader_editor_plugin.h
+++ b/editor/plugins/visual_shader_editor_plugin.h
@@ -493,6 +493,8 @@ class VisualShaderEditor : public VBoxContainer {
 	void _varying_unselected();
 	void _update_varying_tree();
 
+	void _set_custom_node_option(int p_index, int p_node, int p_op);
+
 	Vector2 menu_point;
 	void _node_menu_id_pressed(int p_idx);
 
diff --git a/scene/resources/visual_shader.cpp b/scene/resources/visual_shader.cpp
index ea1207605c..489b866e70 100644
--- a/scene/resources/visual_shader.cpp
+++ b/scene/resources/visual_shader.cpp
@@ -435,8 +435,62 @@ VisualShaderNode::VisualShaderNode() {
 
 /////////////////////////////////////////////////////////
 
+void VisualShaderNodeCustom::update_property_default_values() {
+	int prop_count;
+	if (GDVIRTUAL_CALL(_get_property_count, prop_count)) {
+		for (int i = 0; i < prop_count; i++) {
+			int selected = 0;
+			if (GDVIRTUAL_CALL(_get_property_default_index, i, selected)) {
+				dp_selected_cache[i] = selected;
+			}
+		}
+	}
+}
+
+void VisualShaderNodeCustom::update_input_port_default_values() {
+	int input_port_count;
+	if (GDVIRTUAL_CALL(_get_input_port_count, input_port_count)) {
+		for (int i = 0; i < input_port_count; i++) {
+			Variant value;
+			if (GDVIRTUAL_CALL(_get_input_port_default_value, i, value)) {
+				default_input_values[i] = value;
+			}
+		}
+	}
+}
+
 void VisualShaderNodeCustom::update_ports() {
 	{
+		dp_props.clear();
+		int prop_count;
+		if (GDVIRTUAL_CALL(_get_property_count, prop_count)) {
+			for (int i = 0; i < prop_count; i++) {
+				DropDownListProperty prop;
+				if (!GDVIRTUAL_CALL(_get_property_name, i, prop.name)) {
+					prop.name = "prop";
+				}
+				if (!GDVIRTUAL_CALL(_get_property_options, i, prop.options)) {
+					prop.options.push_back("Default");
+				}
+				dp_props.push_back(prop);
+			}
+		}
+	}
+
+	{
+		Vector<String> vprops = properties.split(";", false);
+		for (int i = 0; i < vprops.size(); i++) {
+			Vector<String> arr = vprops[i].split(",", false);
+			ERR_FAIL_COND(arr.size() != 2);
+			ERR_FAIL_COND(!arr[0].is_valid_int());
+			ERR_FAIL_COND(!arr[1].is_valid_int());
+			int index = arr[0].to_int();
+			int selected = arr[1].to_int();
+			dp_selected_cache[index] = selected;
+		}
+	}
+
+	{
 		input_ports.clear();
 		int input_port_count;
 		if (GDVIRTUAL_CALL(_get_input_port_count, input_port_count)) {
@@ -479,6 +533,15 @@ void VisualShaderNodeCustom::update_ports() {
 	}
 }
 
+void VisualShaderNodeCustom::update_properties() {
+	properties = "";
+	for (const KeyValue<int, int> &p : dp_selected_cache) {
+		if (p.value != 0) {
+			properties += itos(p.key) + "," + itos(p.value) + ";";
+		}
+	}
+}
+
 String VisualShaderNodeCustom::get_caption() const {
 	String ret = "Unnamed";
 	GDVIRTUAL_CALL(_get_name, ret);
@@ -635,6 +698,14 @@ void VisualShaderNodeCustom::_set_initialized(bool p_enabled) {
 	is_initialized = p_enabled;
 }
 
+void VisualShaderNodeCustom::_set_properties(const String &p_properties) {
+	properties = p_properties;
+}
+
+String VisualShaderNodeCustom::_get_properties() const {
+	return properties;
+}
+
 String VisualShaderNodeCustom::_get_name() const {
 	String ret;
 	GDVIRTUAL_CALL(_get_name, ret);
@@ -665,6 +736,21 @@ bool VisualShaderNodeCustom::_is_highend() const {
 	return ret;
 }
 
+void VisualShaderNodeCustom::_set_option_index(int p_option, int p_value) {
+	dp_selected_cache[p_option] = p_value;
+	update_properties();
+	update_ports();
+	update_input_port_default_values();
+	emit_changed();
+}
+
+int VisualShaderNodeCustom::get_option_index(int p_option) const {
+	if (!dp_selected_cache.has(p_option)) {
+		return 0;
+	}
+	return dp_selected_cache[p_option];
+}
+
 void VisualShaderNodeCustom::_bind_methods() {
 	GDVIRTUAL_BIND(_get_name);
 	GDVIRTUAL_BIND(_get_description);
@@ -673,10 +759,15 @@ void VisualShaderNodeCustom::_bind_methods() {
 	GDVIRTUAL_BIND(_get_input_port_count);
 	GDVIRTUAL_BIND(_get_input_port_type, "port");
 	GDVIRTUAL_BIND(_get_input_port_name, "port");
+	GDVIRTUAL_BIND(_get_input_port_default_value, "port");
 	GDVIRTUAL_BIND(_get_default_input_port, "type");
 	GDVIRTUAL_BIND(_get_output_port_count);
 	GDVIRTUAL_BIND(_get_output_port_type, "port");
 	GDVIRTUAL_BIND(_get_output_port_name, "port");
+	GDVIRTUAL_BIND(_get_property_count);
+	GDVIRTUAL_BIND(_get_property_name, "index");
+	GDVIRTUAL_BIND(_get_property_default_index, "index");
+	GDVIRTUAL_BIND(_get_property_options, "index");
 	GDVIRTUAL_BIND(_get_code, "input_vars", "output_vars", "mode", "type");
 	GDVIRTUAL_BIND(_get_func_code, "mode", "type");
 	GDVIRTUAL_BIND(_get_global_code, "mode");
@@ -686,8 +777,14 @@ void VisualShaderNodeCustom::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("_set_initialized", "enabled"), &VisualShaderNodeCustom::_set_initialized);
 	ClassDB::bind_method(D_METHOD("_is_initialized"), &VisualShaderNodeCustom::_is_initialized);
 	ClassDB::bind_method(D_METHOD("_set_input_port_default_value", "port", "value"), &VisualShaderNodeCustom::_set_input_port_default_value);
+	ClassDB::bind_method(D_METHOD("_set_option_index", "option", "value"), &VisualShaderNodeCustom::_set_option_index);
+	ClassDB::bind_method(D_METHOD("_set_properties", "properties"), &VisualShaderNodeCustom::_set_properties);
+	ClassDB::bind_method(D_METHOD("_get_properties"), &VisualShaderNodeCustom::_get_properties);
+
+	ClassDB::bind_method(D_METHOD("get_option_index", "option"), &VisualShaderNodeCustom::get_option_index);
 
 	ADD_PROPERTY(PropertyInfo(Variant::BOOL, "initialized", PROPERTY_HINT_NONE, "", PROPERTY_USAGE_NO_EDITOR | PROPERTY_USAGE_INTERNAL), "_set_initialized", "_is_initialized");
+	ADD_PROPERTY(PropertyInfo(Variant::STRING, "properties", PROPERTY_HINT_NONE, "", PROPERTY_USAGE_NO_EDITOR | PROPERTY_USAGE_INTERNAL), "_set_properties", "_get_properties");
 }
 
 VisualShaderNodeCustom::VisualShaderNodeCustom() {
diff --git a/scene/resources/visual_shader.h b/scene/resources/visual_shader.h
index 1d23b80839..501a538c86 100644
--- a/scene/resources/visual_shader.h
+++ b/scene/resources/visual_shader.h
@@ -257,12 +257,12 @@ class VisualShaderNode : public Resource {
 
 	int port_preview = -1;
 
-	HashMap<int, Variant> default_input_values;
 	HashMap<int, bool> connected_input_ports;
 	HashMap<int, int> connected_output_ports;
 	HashMap<int, bool> expanded_output_ports;
 
 protected:
+	HashMap<int, Variant> default_input_values;
 	bool simple_decl = true;
 	bool disabled = false;
 	bool closable = false;
@@ -363,8 +363,19 @@ class VisualShaderNodeCustom : public VisualShaderNode {
 	bool is_initialized = false;
 	List<Port> input_ports;
 	List<Port> output_ports;
+	struct Property {
+		String name;
+	};
+	struct DropDownListProperty : public Property {
+		Vector<String> options;
+	};
+	HashMap<int, int> dp_selected_cache;
+	HashMap<int, int> dp_default_cache;
+	List<DropDownListProperty> dp_props;
+	String properties;
 
 	friend class VisualShaderEditor;
+	friend class VisualShaderGraphPlugin;
 
 protected:
 	virtual String get_caption() const override;
@@ -390,10 +401,15 @@ protected:
 	GDVIRTUAL0RC(int, _get_input_port_count)
 	GDVIRTUAL1RC(PortType, _get_input_port_type, int)
 	GDVIRTUAL1RC(String, _get_input_port_name, int)
+	GDVIRTUAL1RC(Variant, _get_input_port_default_value, int)
 	GDVIRTUAL1RC(int, _get_default_input_port, PortType)
 	GDVIRTUAL0RC(int, _get_output_port_count)
 	GDVIRTUAL1RC(PortType, _get_output_port_type, int)
 	GDVIRTUAL1RC(String, _get_output_port_name, int)
+	GDVIRTUAL0RC(int, _get_property_count)
+	GDVIRTUAL1RC(String, _get_property_name, int)
+	GDVIRTUAL1RC(int, _get_property_default_index, int)
+	GDVIRTUAL1RC(Vector<String>, _get_property_options, 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)
@@ -414,16 +430,24 @@ protected:
 
 public:
 	VisualShaderNodeCustom();
+	void update_property_default_values();
+	void update_input_port_default_values();
 	void update_ports();
+	void update_properties();
 
 	bool _is_initialized();
 	void _set_initialized(bool p_enabled);
+	void _set_properties(const String &p_properties);
+	String _get_properties() const;
 
 	String _get_name() const;
 	String _get_description() const;
 	String _get_category() const;
 	PortType _get_return_icon_type() const;
 	bool _is_highend() const;
+	void _set_option_index(int p_op, int p_index);
+
+	int get_option_index(int p_op) const;
 };
 
 /////