Added "Is" and "Compare" functions to visual shaders

3 files changed, 427 insertions, 0 deletions

diff --git a/scene/register_scene_types.cpp b/scene/register_scene_types.cpp
index 0423fcb5f0..9a21da7479 100644
--- a/scene/register_scene_types.cpp
+++ b/scene/register_scene_types.cpp
@@ -527,6 +527,8 @@ void register_scene_types() {
 	ClassDB::register_class<VisualShaderNodeSwitch>();
 	ClassDB::register_class<VisualShaderNodeFresnel>();
 	ClassDB::register_class<VisualShaderNodeExpression>();
+	ClassDB::register_class<VisualShaderNodeIs>();
+	ClassDB::register_class<VisualShaderNodeCompare>();
 
 	ClassDB::register_class<ShaderMaterial>();
 	ClassDB::register_virtual_class<CanvasItem>();
diff --git a/scene/resources/visual_shader_nodes.cpp b/scene/resources/visual_shader_nodes.cpp
index 746edc65b0..235b56b0ce 100644
--- a/scene/resources/visual_shader_nodes.cpp
+++ b/scene/resources/visual_shader_nodes.cpp
@@ -3200,3 +3200,317 @@ VisualShaderNodeFresnel::VisualShaderNodeFresnel() {
 	set_input_port_default_value(2, false);
 	set_input_port_default_value(3, 1.0);
 }
+
+////////////// Is
+
+String VisualShaderNodeIs::get_caption() const {
+
+	return "Is";
+}
+
+int VisualShaderNodeIs::get_input_port_count() const {
+
+	return 1;
+}
+
+VisualShaderNodeIs::PortType VisualShaderNodeIs::get_input_port_type(int p_port) const {
+
+	return PORT_TYPE_SCALAR;
+}
+
+String VisualShaderNodeIs::get_input_port_name(int p_port) const {
+
+	return "";
+}
+
+int VisualShaderNodeIs::get_output_port_count() const {
+
+	return 1;
+}
+
+VisualShaderNodeIs::PortType VisualShaderNodeIs::get_output_port_type(int p_port) const {
+
+	return PORT_TYPE_BOOLEAN;
+}
+
+String VisualShaderNodeIs::get_output_port_name(int p_port) const {
+
+	return "";
+}
+
+String VisualShaderNodeIs::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 {
+
+	static const char *funcs[FUNC_IS_NAN + 1] = {
+		"isinf($)",
+		"isnan($)"
+	};
+
+	String code;
+	code += "\t" + p_output_vars[0] + "=" + String(funcs[func]).replace("$", p_input_vars[0]) + ";\n";
+	return code;
+}
+
+void VisualShaderNodeIs::set_function(Function p_func) {
+
+	func = p_func;
+	emit_changed();
+}
+
+VisualShaderNodeIs::Function VisualShaderNodeIs::get_function() const {
+
+	return func;
+}
+
+Vector<StringName> VisualShaderNodeIs::get_editable_properties() const {
+
+	Vector<StringName> props;
+	props.push_back("function");
+	return props;
+}
+
+void VisualShaderNodeIs::_bind_methods() {
+
+	ClassDB::bind_method(D_METHOD("set_function", "func"), &VisualShaderNodeIs::set_function);
+	ClassDB::bind_method(D_METHOD("get_function"), &VisualShaderNodeIs::get_function);
+
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "function", PROPERTY_HINT_ENUM, "Inf,NaN"), "set_function", "get_function");
+
+	BIND_ENUM_CONSTANT(FUNC_IS_INF);
+	BIND_ENUM_CONSTANT(FUNC_IS_NAN);
+}
+
+VisualShaderNodeIs::VisualShaderNodeIs() {
+
+	func = FUNC_IS_INF;
+	set_input_port_default_value(0, 0.0);
+}
+
+////////////// Compare
+
+String VisualShaderNodeCompare::get_caption() const {
+
+	return "Compare";
+}
+
+int VisualShaderNodeCompare::get_input_port_count() const {
+
+	if (ctype == CTYPE_SCALAR && (func == FUNC_EQUAL || func == FUNC_NOT_EQUAL)) {
+		return 3;
+	}
+	return 2;
+}
+
+VisualShaderNodeCompare::PortType VisualShaderNodeCompare::get_input_port_type(int p_port) const {
+
+	if (p_port == 2)
+		return PORT_TYPE_SCALAR;
+	switch (ctype) {
+		case CTYPE_SCALAR:
+			return PORT_TYPE_SCALAR;
+		case CTYPE_VECTOR:
+			return PORT_TYPE_VECTOR;
+		case CTYPE_BOOLEAN:
+			return PORT_TYPE_BOOLEAN;
+		case CTYPE_TRANSFORM:
+			return PORT_TYPE_TRANSFORM;
+	}
+	return PORT_TYPE_VECTOR;
+}
+
+String VisualShaderNodeCompare::get_input_port_name(int p_port) const {
+	if (p_port == 0)
+		return "a";
+	else if (p_port == 1)
+		return "b";
+	else if (p_port == 2)
+		return "tolerance";
+	return "";
+}
+
+int VisualShaderNodeCompare::get_output_port_count() const {
+	return 1;
+}
+
+VisualShaderNodeCompare::PortType VisualShaderNodeCompare::get_output_port_type(int p_port) const {
+	return PORT_TYPE_BOOLEAN;
+}
+
+String VisualShaderNodeCompare::get_output_port_name(int p_port) const {
+	if (p_port == 0)
+		return "result";
+	return "";
+}
+
+String VisualShaderNodeCompare::get_warning(Shader::Mode p_mode, VisualShader::Type p_type) const {
+
+	if (ctype == CTYPE_BOOLEAN || ctype == CTYPE_TRANSFORM) {
+		if (func > FUNC_NOT_EQUAL) {
+			return TTR("Invalid comparsion function for that type.");
+		}
+	}
+
+	return "";
+}
+
+String VisualShaderNodeCompare::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 {
+
+	static const char *ops[FUNC_LESS_THAN_EQUAL + 1] = {
+		"==",
+		"!=",
+		">",
+		">=",
+		"<",
+		"<=",
+	};
+
+	static const char *funcs[FUNC_LESS_THAN_EQUAL + 1] = {
+		"equal($)",
+		"notEqual($)",
+		"greaterThan($)",
+		"greaterThanEqual($)",
+		"lessThan($)",
+		"lessThanEqual($)",
+	};
+
+	static const char *conds[COND_ANY + 1] = {
+		"all($)",
+		"any($)",
+	};
+
+	String code;
+	switch (ctype) {
+		case CTYPE_SCALAR:
+			if (func == FUNC_EQUAL) {
+				code += "\t" + p_output_vars[0] + "=(abs(" + p_input_vars[0] + "-" + p_input_vars[1] + ")<" + p_input_vars[2] + ");";
+			} else if (func == FUNC_NOT_EQUAL) {
+				code += "\t" + p_output_vars[0] + "=!(abs(" + p_input_vars[0] + "-" + p_input_vars[1] + ")<" + p_input_vars[2] + ");";
+			} else {
+				code += "\t" + p_output_vars[0] + "=" + (p_input_vars[0] + "$" + p_input_vars[1]).replace("$", ops[func]) + ";\n";
+			}
+			break;
+
+		case CTYPE_VECTOR:
+			code += "\t{\n";
+			code += "\t\tbvec3 _bv=" + String(funcs[func]).replace("$", p_input_vars[0] + ", " + p_input_vars[1]) + ";\n";
+			code += "\t\t" + p_output_vars[0] + "=" + String(conds[condition]).replace("$", "_bv") + ";\n";
+			code += "\t}\n";
+			break;
+
+		case CTYPE_BOOLEAN:
+			if (func > FUNC_NOT_EQUAL)
+				return "\t" + p_output_vars[0] + "=false;\n";
+			code += "\t" + p_output_vars[0] + "=" + (p_input_vars[0] + "$" + p_input_vars[1]).replace("$", ops[func]) + ";\n";
+			break;
+
+		case CTYPE_TRANSFORM:
+			if (func > FUNC_NOT_EQUAL)
+				return "\t" + p_output_vars[0] + "=false;\n";
+			code += "\t" + p_output_vars[0] + "=" + (p_input_vars[0] + "$" + p_input_vars[1]).replace("$", ops[func]) + ";\n";
+			break;
+
+		default:
+			break;
+	}
+	return code;
+}
+
+void VisualShaderNodeCompare::set_comparsion_type(ComparsionType p_type) {
+
+	ctype = p_type;
+
+	switch (ctype) {
+		case CTYPE_SCALAR:
+			set_input_port_default_value(0, 0.0);
+			set_input_port_default_value(1, 0.0);
+			break;
+		case CTYPE_VECTOR:
+			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));
+			break;
+		case CTYPE_BOOLEAN:
+			set_input_port_default_value(0, false);
+			set_input_port_default_value(1, false);
+			break;
+		case CTYPE_TRANSFORM:
+			set_input_port_default_value(0, Transform());
+			set_input_port_default_value(1, Transform());
+			break;
+	}
+	emit_changed();
+}
+
+VisualShaderNodeCompare::ComparsionType VisualShaderNodeCompare::get_comparsion_type() const {
+
+	return ctype;
+}
+
+void VisualShaderNodeCompare::set_function(Function p_func) {
+
+	func = p_func;
+	emit_changed();
+}
+
+VisualShaderNodeCompare::Function VisualShaderNodeCompare::get_function() const {
+
+	return func;
+}
+
+void VisualShaderNodeCompare::set_condition(Condition p_cond) {
+
+	condition = p_cond;
+	emit_changed();
+}
+
+VisualShaderNodeCompare::Condition VisualShaderNodeCompare::get_condition() const {
+
+	return condition;
+}
+
+Vector<StringName> VisualShaderNodeCompare::get_editable_properties() const {
+	Vector<StringName> props;
+	props.push_back("type");
+	props.push_back("function");
+	if (ctype == CTYPE_VECTOR)
+		props.push_back("condition");
+	return props;
+}
+
+void VisualShaderNodeCompare::_bind_methods() {
+
+	ClassDB::bind_method(D_METHOD("set_comparsion_type", "type"), &VisualShaderNodeCompare::set_comparsion_type);
+	ClassDB::bind_method(D_METHOD("get_comparsion_type"), &VisualShaderNodeCompare::get_comparsion_type);
+
+	ClassDB::bind_method(D_METHOD("set_function", "func"), &VisualShaderNodeCompare::set_function);
+	ClassDB::bind_method(D_METHOD("get_function"), &VisualShaderNodeCompare::get_function);
+
+	ClassDB::bind_method(D_METHOD("set_condition", "condition"), &VisualShaderNodeCompare::set_condition);
+	ClassDB::bind_method(D_METHOD("get_condition"), &VisualShaderNodeCompare::get_condition);
+
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "type", PROPERTY_HINT_ENUM, "Scalar,Vector,Boolean,Transform"), "set_comparsion_type", "get_comparsion_type");
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "function", PROPERTY_HINT_ENUM, "a == b,a != b,a > b,a >= b,a < b,a <= b"), "set_function", "get_function");
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "condition", PROPERTY_HINT_ENUM, "All,Any"), "set_condition", "get_condition");
+
+	BIND_ENUM_CONSTANT(CTYPE_SCALAR);
+	BIND_ENUM_CONSTANT(CTYPE_VECTOR);
+	BIND_ENUM_CONSTANT(CTYPE_BOOLEAN);
+	BIND_ENUM_CONSTANT(CTYPE_TRANSFORM);
+
+	BIND_ENUM_CONSTANT(FUNC_EQUAL);
+	BIND_ENUM_CONSTANT(FUNC_NOT_EQUAL);
+	BIND_ENUM_CONSTANT(FUNC_GREATER_THAN);
+	BIND_ENUM_CONSTANT(FUNC_GREATER_THAN_EQUAL);
+	BIND_ENUM_CONSTANT(FUNC_LESS_THAN);
+	BIND_ENUM_CONSTANT(FUNC_LESS_THAN_EQUAL);
+
+	BIND_ENUM_CONSTANT(COND_ALL);
+	BIND_ENUM_CONSTANT(COND_ANY);
+}
+
+VisualShaderNodeCompare::VisualShaderNodeCompare() {
+	ctype = CTYPE_SCALAR;
+	func = FUNC_EQUAL;
+	condition = COND_ALL;
+	set_input_port_default_value(0, 0.0);
+	set_input_port_default_value(1, 0.0);
+	set_input_port_default_value(2, CMP_EPSILON);
+}
diff --git a/scene/resources/visual_shader_nodes.h b/scene/resources/visual_shader_nodes.h
index 235714f697..dca0ebe033 100644
--- a/scene/resources/visual_shader_nodes.h
+++ b/scene/resources/visual_shader_nodes.h
@@ -1513,4 +1513,115 @@ public:
 	VisualShaderNodeFresnel();
 };
 
+///////////////////////////////////////
+/// Is
+///////////////////////////////////////
+
+class VisualShaderNodeIs : public VisualShaderNode {
+	GDCLASS(VisualShaderNodeIs, VisualShaderNode);
+
+public:
+	enum Function {
+		FUNC_IS_INF,
+		FUNC_IS_NAN,
+	};
+
+protected:
+	Function func;
+
+protected:
+	static void _bind_methods();
+
+public:
+	virtual String get_caption() const;
+
+	virtual int get_input_port_count() const;
+	virtual PortType get_input_port_type(int p_port) const;
+	virtual String get_input_port_name(int p_port) const;
+
+	virtual int get_output_port_count() const;
+	virtual PortType get_output_port_type(int p_port) const;
+	virtual String get_output_port_name(int p_port) 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; //if no output is connected, the output var passed will be empty. if no input is connected and input is NIL, the input var passed will be empty
+
+	void set_function(Function p_func);
+	Function get_function() const;
+
+	virtual Vector<StringName> get_editable_properties() const;
+
+	VisualShaderNodeIs();
+};
+
+VARIANT_ENUM_CAST(VisualShaderNodeIs::Function)
+
+///////////////////////////////////////
+/// Compare
+///////////////////////////////////////
+
+class VisualShaderNodeCompare : public VisualShaderNode {
+	GDCLASS(VisualShaderNodeCompare, VisualShaderNode);
+
+public:
+	enum ComparsionType {
+		CTYPE_SCALAR,
+		CTYPE_VECTOR,
+		CTYPE_BOOLEAN,
+		CTYPE_TRANSFORM
+	};
+
+	enum Function {
+		FUNC_EQUAL,
+		FUNC_NOT_EQUAL,
+		FUNC_GREATER_THAN,
+		FUNC_GREATER_THAN_EQUAL,
+		FUNC_LESS_THAN,
+		FUNC_LESS_THAN_EQUAL,
+	};
+
+	enum Condition {
+		COND_ALL,
+		COND_ANY,
+	};
+
+protected:
+	ComparsionType ctype;
+	Function func;
+	Condition condition;
+
+protected:
+	static void _bind_methods();
+
+public:
+	virtual String get_caption() const;
+
+	virtual int get_input_port_count() const;
+	virtual PortType get_input_port_type(int p_port) const;
+	virtual String get_input_port_name(int p_port) const;
+
+	virtual int get_output_port_count() const;
+	virtual PortType get_output_port_type(int p_port) const;
+	virtual String get_output_port_name(int p_port) 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; //if no output is connected, the output var passed will be empty. if no input is connected and input is NIL, the input var passed will be empty
+
+	void set_comparsion_type(ComparsionType p_func);
+	ComparsionType get_comparsion_type() const;
+
+	void set_function(Function p_func);
+	Function get_function() const;
+
+	void set_condition(Condition p_mode);
+	Condition get_condition() const;
+
+	virtual Vector<StringName> get_editable_properties() const;
+	virtual String get_warning(Shader::Mode p_mode, VisualShader::Type p_type) const;
+
+	VisualShaderNodeCompare();
+};
+
+VARIANT_ENUM_CAST(VisualShaderNodeCompare::ComparsionType)
+VARIANT_ENUM_CAST(VisualShaderNodeCompare::Function)
+VARIANT_ENUM_CAST(VisualShaderNodeCompare::Condition)
+
 #endif // VISUAL_SHADER_NODES_H