Allow vec2 and vec4 for reflect and refract

2 files changed, 30 insertions, 24 deletions

diff --git a/scene/resources/visual_shader_nodes.cpp b/scene/resources/visual_shader_nodes.cpp
index b8667f07fe..fd1edfccec 100644
--- a/scene/resources/visual_shader_nodes.cpp
+++ b/scene/resources/visual_shader_nodes.cpp
@@ -1901,13 +1901,7 @@ String VisualShaderNodeVectorOp::generate_code(Shader::Mode p_mode, VisualShader
 			code += "atan(" + p_input_vars[0] + ", " + p_input_vars[1] + ");\n";
 			break;
 		case OP_REFLECT:
-			if (op_type == OP_TYPE_VECTOR_2D) { // Not supported.
-				code += "vec2(0.0);\n";
-			} else if (op_type == OP_TYPE_VECTOR_4D) { // Not supported.
-				code += "vec4(0.0);\n";
-			} else {
-				code += "reflect(" + p_input_vars[0] + ", " + p_input_vars[1] + ");\n";
-			}
+			code += "reflect(" + p_input_vars[0] + ", " + p_input_vars[1] + ");\n";
 			break;
 		case OP_STEP:
 			code += "step(" + p_input_vars[0] + ", " + p_input_vars[1] + ");\n";
@@ -1967,7 +1961,7 @@ String VisualShaderNodeVectorOp::get_warning(Shader::Mode p_mode, VisualShader::
 	bool invalid_type = false;
 
 	if (op_type == OP_TYPE_VECTOR_2D || op_type == OP_TYPE_VECTOR_4D) {
-		if (op == OP_CROSS || op == OP_REFLECT) {
+		if (op == OP_CROSS) {
 			invalid_type = true;
 		}
 	}
@@ -4006,14 +4000,6 @@ int VisualShaderNodeVectorRefract::get_input_port_count() const {
 	return 3;
 }
 
-VisualShaderNodeVectorRefract::PortType VisualShaderNodeVectorRefract::get_input_port_type(int p_port) const {
-	if (p_port == 2) {
-		return PORT_TYPE_SCALAR;
-	}
-
-	return PORT_TYPE_VECTOR_3D;
-}
-
 String VisualShaderNodeVectorRefract::get_input_port_name(int p_port) const {
 	switch (p_port) {
 		case 0:
@@ -4030,10 +4016,6 @@ int VisualShaderNodeVectorRefract::get_output_port_count() const {
 	return 1;
 }
 
-VisualShaderNodeVectorRefract::PortType VisualShaderNodeVectorRefract::get_output_port_type(int p_port) const {
-	return PORT_TYPE_VECTOR_3D;
-}
-
 String VisualShaderNodeVectorRefract::get_output_port_name(int p_port) const {
 	return "";
 }
@@ -4042,6 +4024,31 @@ String VisualShaderNodeVectorRefract::generate_code(Shader::Mode p_mode, VisualS
 	return "	" + p_output_vars[0] + " = refract(" + p_input_vars[0] + ", " + p_input_vars[1] + ", " + p_input_vars[2] + ");\n";
 }
 
+void VisualShaderNodeVectorRefract::set_op_type(OpType p_op_type) {
+	ERR_FAIL_INDEX(int(p_op_type), int(OP_TYPE_MAX));
+	if (op_type == p_op_type) {
+		return;
+	}
+	switch (p_op_type) {
+		case OP_TYPE_VECTOR_2D: {
+			set_input_port_default_value(0, Vector2(), get_input_port_default_value(0));
+			set_input_port_default_value(1, Vector2(), get_input_port_default_value(1));
+		} break;
+		case OP_TYPE_VECTOR_3D: {
+			set_input_port_default_value(0, Vector3(), get_input_port_default_value(0));
+			set_input_port_default_value(1, Vector3(), get_input_port_default_value(1));
+		} break;
+		case OP_TYPE_VECTOR_4D: {
+			set_input_port_default_value(0, Quaternion(), get_input_port_default_value(0));
+			set_input_port_default_value(1, Quaternion(), get_input_port_default_value(1));
+		} break;
+		default:
+			break;
+	}
+	op_type = p_op_type;
+	emit_changed();
+}
+
 VisualShaderNodeVectorRefract::VisualShaderNodeVectorRefract() {
 	set_input_port_default_value(0, Vector3(0.0, 0.0, 0.0));
 	set_input_port_default_value(1, Vector3(0.0, 0.0, 0.0));
diff --git a/scene/resources/visual_shader_nodes.h b/scene/resources/visual_shader_nodes.h
index 1eb7b7240f..4f00c3c2ea 100644
--- a/scene/resources/visual_shader_nodes.h
+++ b/scene/resources/visual_shader_nodes.h
@@ -1564,21 +1564,20 @@ public:
 /// REFRACT
 ///////////////////////////////////////
 
-class VisualShaderNodeVectorRefract : public VisualShaderNode {
-	GDCLASS(VisualShaderNodeVectorRefract, VisualShaderNode);
+class VisualShaderNodeVectorRefract : public VisualShaderNodeVectorBase {
+	GDCLASS(VisualShaderNodeVectorRefract, VisualShaderNodeVectorBase);
 
 public:
 	virtual String get_caption() const override;
 
 	virtual int get_input_port_count() const override;
-	virtual PortType get_input_port_type(int p_port) const override;
 	virtual String get_input_port_name(int p_port) const override;
 
 	virtual int get_output_port_count() const override;
-	virtual PortType get_output_port_type(int p_port) const override;
 	virtual String get_output_port_name(int p_port) const override;
 
 	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 void set_op_type(OpType p_op_type) override;
 
 	VisualShaderNodeVectorRefract();
 };