Some improvements for Switch node in visual shaders

4 files changed, 48 insertions, 3 deletions

diff --git a/editor/plugins/visual_shader_editor_plugin.cpp b/editor/plugins/visual_shader_editor_plugin.cpp
index 6b338ca02b..91a26d5da3 100644
--- a/editor/plugins/visual_shader_editor_plugin.cpp
+++ b/editor/plugins/visual_shader_editor_plugin.cpp
@@ -2367,6 +2367,7 @@ VisualShaderEditor::VisualShaderEditor() {
 	add_options.push_back(AddOption("LessThanEqual", "Conditional", "Functions", "VisualShaderNodeCompare", vformat(compare_func_desc, TTR("Less Than or Equal (<=)")), VisualShaderNodeCompare::FUNC_LESS_THAN_EQUAL, VisualShaderNode::PORT_TYPE_BOOLEAN));
 	add_options.push_back(AddOption("NotEqual", "Conditional", "Functions", "VisualShaderNodeCompare", vformat(compare_func_desc, TTR("Not Equal (!=)")), VisualShaderNodeCompare::FUNC_NOT_EQUAL, VisualShaderNode::PORT_TYPE_BOOLEAN));
 	add_options.push_back(AddOption("Switch", "Conditional", "Functions", "VisualShaderNodeSwitch", TTR("Returns an associated vector if the provided boolean value is true or false."), -1, VisualShaderNode::PORT_TYPE_VECTOR));
+	add_options.push_back(AddOption("SwitchS", "Conditional", "Functions", "VisualShaderNodeScalarSwitch", TTR("Returns an associated scalar if the provided boolean value is true or false."), -1, VisualShaderNode::PORT_TYPE_SCALAR));
 
 	add_options.push_back(AddOption("Compare", "Conditional", "Common", "VisualShaderNodeCompare", TTR("Returns the boolean result of the comparison between two parameters."), -1, VisualShaderNode::PORT_TYPE_BOOLEAN));
 	add_options.push_back(AddOption("Is", "Conditional", "Common", "VisualShaderNodeIs", TTR("Returns the boolean result of the comparison between INF (or NaN) and a scalar parameter."), -1, VisualShaderNode::PORT_TYPE_BOOLEAN));
diff --git a/scene/register_scene_types.cpp b/scene/register_scene_types.cpp
index 06d84302a3..5b21092f60 100644
--- a/scene/register_scene_types.cpp
+++ b/scene/register_scene_types.cpp
@@ -528,6 +528,7 @@ void register_scene_types() {
 	ClassDB::register_class<VisualShaderNodeCubeMapUniform>();
 	ClassDB::register_class<VisualShaderNodeIf>();
 	ClassDB::register_class<VisualShaderNodeSwitch>();
+	ClassDB::register_class<VisualShaderNodeScalarSwitch>();
 	ClassDB::register_class<VisualShaderNodeFresnel>();
 	ClassDB::register_class<VisualShaderNodeExpression>();
 	ClassDB::register_class<VisualShaderNodeGlobalExpression>();
diff --git a/scene/resources/visual_shader_nodes.cpp b/scene/resources/visual_shader_nodes.cpp
index 9195d80cfc..e4e0414342 100644
--- a/scene/resources/visual_shader_nodes.cpp
+++ b/scene/resources/visual_shader_nodes.cpp
@@ -3227,7 +3227,7 @@ VisualShaderNodeIf::VisualShaderNodeIf() {
 ////////////// Switch
 
 String VisualShaderNodeSwitch::get_caption() const {
-	return "Switch";
+	return "VectorSwitch";
 }
 
 int VisualShaderNodeSwitch::get_input_port_count() const {
@@ -3269,7 +3269,15 @@ String VisualShaderNodeSwitch::get_output_port_name(int p_port) const {
 String VisualShaderNodeSwitch::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) const {
 
 	String code;
-	code += "\tif(" + p_input_vars[0] + ")\n";
+	if (p_input_vars[0] == String()) {
+		if ((bool)get_input_port_default_value(0)) {
+			code += "\tif(true)\n";
+		} else {
+			code += "\tif(false)\n";
+		}
+	} else {
+		code += "\tif(" + p_input_vars[0] + ")\n";
+	}
 	code += "\t{\n";
 	code += "\t\t" + p_output_vars[0] + "=" + p_input_vars[1] + ";\n";
 	code += "\t}\n";
@@ -3282,10 +3290,33 @@ String VisualShaderNodeSwitch::generate_code(Shader::Mode p_mode, VisualShader::
 
 VisualShaderNodeSwitch::VisualShaderNodeSwitch() {
 	set_input_port_default_value(0, false);
-	set_input_port_default_value(1, Vector3(0.0, 0.0, 0.0));
+	set_input_port_default_value(1, Vector3(1.0, 1.0, 1.0));
 	set_input_port_default_value(2, Vector3(0.0, 0.0, 0.0));
 }
 
+////////////// Switch(scalar)
+
+String VisualShaderNodeScalarSwitch::get_caption() const {
+	return "ScalarSwitch";
+}
+
+VisualShaderNodeScalarSwitch::PortType VisualShaderNodeScalarSwitch::get_input_port_type(int p_port) const {
+	if (p_port == 0) {
+		return PORT_TYPE_BOOLEAN;
+	}
+	return PORT_TYPE_SCALAR;
+}
+
+VisualShaderNodeScalarSwitch::PortType VisualShaderNodeScalarSwitch::get_output_port_type(int p_port) const {
+	return PORT_TYPE_SCALAR;
+}
+
+VisualShaderNodeScalarSwitch::VisualShaderNodeScalarSwitch() {
+	set_input_port_default_value(0, false);
+	set_input_port_default_value(1, 1.0);
+	set_input_port_default_value(2, 0.0);
+}
+
 ////////////// Fresnel
 
 String VisualShaderNodeFresnel::get_caption() const {
diff --git a/scene/resources/visual_shader_nodes.h b/scene/resources/visual_shader_nodes.h
index f7efa396dc..0c6e060353 100644
--- a/scene/resources/visual_shader_nodes.h
+++ b/scene/resources/visual_shader_nodes.h
@@ -1530,6 +1530,18 @@ public:
 	VisualShaderNodeSwitch();
 };
 
+class VisualShaderNodeScalarSwitch : public VisualShaderNodeSwitch {
+	GDCLASS(VisualShaderNodeScalarSwitch, VisualShaderNodeSwitch);
+
+public:
+	virtual String get_caption() const;
+
+	virtual PortType get_input_port_type(int p_port) const;
+	virtual PortType get_output_port_type(int p_port) const;
+
+	VisualShaderNodeScalarSwitch();
+};
+
 ///////////////////////////////////////
 /// FRESNEL
 ///////////////////////////////////////