Allow passing varying from fragment to light shader function

4 files changed, 175 insertions, 26 deletions

diff --git a/servers/rendering/renderer_rd/shader_compiler_rd.cpp b/servers/rendering/renderer_rd/shader_compiler_rd.cpp
index e77141b26c..a81706700a 100644
--- a/servers/rendering/renderer_rd/shader_compiler_rd.cpp
+++ b/servers/rendering/renderer_rd/shader_compiler_rd.cpp
@@ -687,7 +687,15 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 
 			uint32_t index = p_default_actions.base_varying_index;
 
+			List<Pair<StringName, SL::ShaderNode::Varying>> var_frag_to_light;
+
 			for (Map<StringName, SL::ShaderNode::Varying>::Element *E = pnode->varyings.front(); E; E = E->next()) {
+				if (E->get().stage == SL::ShaderNode::Varying::STAGE_FRAGMENT_TO_LIGHT || E->get().stage == SL::ShaderNode::Varying::STAGE_FRAGMENT) {
+					var_frag_to_light.push_back(Pair<StringName, SL::ShaderNode::Varying>(E->key(), E->get()));
+					fragment_varyings.insert(E->key());
+					continue;
+				}
+
 				String vcode;
 				String interp_mode = _interpstr(E->get().interpolation);
 				vcode += _prestr(E->get().precision);
@@ -705,6 +713,21 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 				index++;
 			}
 
+			if (var_frag_to_light.size() > 0) {
+				String gcode = "\n\nstruct {\n";
+				for (List<Pair<StringName, SL::ShaderNode::Varying>>::Element *E = var_frag_to_light.front(); E; E = E->next()) {
+					gcode += "\t" + _prestr(E->get().second.precision) + _typestr(E->get().second.type) + " " + _mkid(E->get().first);
+					if (E->get().second.array_size > 0) {
+						gcode += "[";
+						gcode += itos(E->get().second.array_size);
+						gcode += "]";
+					}
+					gcode += ";\n";
+				}
+				gcode += "} frag_to_light;\n";
+				r_gen_code.fragment_global += gcode;
+			}
+
 			for (int i = 0; i < pnode->vconstants.size(); i++) {
 				const SL::ShaderNode::Constant &cnode = pnode->vconstants[i];
 				String gcode;
@@ -833,6 +856,19 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 		} break;
 		case SL::Node::TYPE_VARIABLE: {
 			SL::VariableNode *vnode = (SL::VariableNode *)p_node;
+			bool use_fragment_varying = false;
+
+			if (current_func_name != vertex_name) {
+				if (p_assigning) {
+					if (shader->varyings.has(vnode->name)) {
+						use_fragment_varying = true;
+					}
+				} else {
+					if (fragment_varyings.has(vnode->name)) {
+						use_fragment_varying = true;
+					}
+				}
+			}
 
 			if (p_assigning && p_actions.write_flag_pointers.has(vnode->name)) {
 				*p_actions.write_flag_pointers[vnode->name] = true;
@@ -877,7 +913,10 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 					}
 
 				} else {
-					code = _mkid(vnode->name); //its something else (local var most likely) use as is
+					if (use_fragment_varying) {
+						code = "frag_to_light.";
+					}
+					code += _mkid(vnode->name); //its something else (local var most likely) use as is
 				}
 			}
 
@@ -962,6 +1001,23 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 		} break;
 		case SL::Node::TYPE_ARRAY: {
 			SL::ArrayNode *anode = (SL::ArrayNode *)p_node;
+			bool use_fragment_varying = false;
+
+			if (current_func_name != vertex_name) {
+				if (anode->assign_expression != nullptr) {
+					use_fragment_varying = true;
+				} else {
+					if (p_assigning) {
+						if (shader->varyings.has(anode->name)) {
+							use_fragment_varying = true;
+						}
+					} else {
+						if (fragment_varyings.has(anode->name)) {
+							use_fragment_varying = true;
+						}
+					}
+				}
+			}
 
 			if (p_assigning && p_actions.write_flag_pointers.has(anode->name)) {
 				*p_actions.write_flag_pointers[anode->name] = true;
@@ -984,7 +1040,10 @@ String ShaderCompilerRD::_dump_node_code(const SL::Node *p_node, int p_level, Ge
 			if (p_default_actions.renames.has(anode->name)) {
 				code = p_default_actions.renames[anode->name];
 			} else {
-				code = _mkid(anode->name);
+				if (use_fragment_varying) {
+					code = "frag_to_light.";
+				}
+				code += _mkid(anode->name);
 			}
 
 			if (anode->call_expression != nullptr) {
@@ -1277,6 +1336,7 @@ Error ShaderCompilerRD::compile(RS::ShaderMode p_mode, const String &p_code, Ide
 	used_name_defines.clear();
 	used_rmode_defines.clear();
 	used_flag_pointers.clear();
+	fragment_varyings.clear();
 
 	shader = parser.get_shader();
 	function = nullptr;
diff --git a/servers/rendering/renderer_rd/shader_compiler_rd.h b/servers/rendering/renderer_rd/shader_compiler_rd.h
index d127d8e01c..6575829e73 100644
--- a/servers/rendering/renderer_rd/shader_compiler_rd.h
+++ b/servers/rendering/renderer_rd/shader_compiler_rd.h
@@ -114,6 +114,7 @@ private:
 	Set<StringName> used_flag_pointers;
 	Set<StringName> used_rmode_defines;
 	Set<StringName> internal_functions;
+	Set<StringName> fragment_varyings;
 
 	DefaultIdentifierActions actions;
 
diff --git a/servers/rendering/shader_language.cpp b/servers/rendering/shader_language.cpp
index e6415c0258..e52d97e4d9 100644
--- a/servers/rendering/shader_language.cpp
+++ b/servers/rendering/shader_language.cpp
@@ -3102,6 +3102,72 @@ bool ShaderLanguage::_is_operator_assign(Operator p_op) const {
 	return false;
 }
 
+bool ShaderLanguage::_validate_varying_assign(ShaderNode::Varying &p_varying, String *r_message) {
+	if (current_function == String("light")) {
+		*r_message = RTR("Varying may not be assigned in the 'light' function.");
+		return false;
+	}
+	switch (p_varying.stage) {
+		case ShaderNode::Varying::STAGE_UNKNOWN: // first assign
+			if (current_function == String("vertex")) {
+				p_varying.stage = ShaderNode::Varying::STAGE_VERTEX;
+			} else if (current_function == String("fragment")) {
+				p_varying.stage = ShaderNode::Varying::STAGE_FRAGMENT;
+			}
+			break;
+		case ShaderNode::Varying::STAGE_VERTEX:
+			if (current_function == String("fragment")) {
+				*r_message = RTR("Varyings which assigned in 'vertex' function may not be reassigned in 'fragment' or 'light'.");
+				return false;
+			}
+			break;
+		case ShaderNode::Varying::STAGE_FRAGMENT:
+			if (current_function == String("vertex")) {
+				*r_message = RTR("Varyings which assigned in 'fragment' function may not be reassigned in 'vertex' or 'light'.");
+				return false;
+			}
+			break;
+		default:
+			break;
+	}
+	return true;
+}
+
+bool ShaderLanguage::_validate_varying_using(ShaderNode::Varying &p_varying, String *r_message) {
+	switch (p_varying.stage) {
+		case ShaderNode::Varying::STAGE_UNKNOWN:
+			*r_message = RTR("Varying must be assigned before using!");
+			return false;
+		case ShaderNode::Varying::STAGE_VERTEX:
+			if (current_function == String("fragment")) {
+				p_varying.stage = ShaderNode::Varying::STAGE_VERTEX_TO_FRAGMENT;
+			} else if (current_function == String("light")) {
+				p_varying.stage = ShaderNode::Varying::STAGE_VERTEX_TO_LIGHT;
+			}
+			break;
+		case ShaderNode::Varying::STAGE_FRAGMENT:
+			if (current_function == String("light")) {
+				p_varying.stage = ShaderNode::Varying::STAGE_FRAGMENT_TO_LIGHT;
+			}
+			break;
+		case ShaderNode::Varying::STAGE_VERTEX_TO_FRAGMENT:
+			if (current_function == String("light")) {
+				*r_message = RTR("Varying must only be used in two different stages, which can be 'vertex' 'fragment' and 'light'");
+				return false;
+			}
+			break;
+		case ShaderNode::Varying::STAGE_VERTEX_TO_LIGHT:
+			if (current_function == String("fragment")) {
+				*r_message = RTR("Varying must only be used in two different stages, which can be 'vertex' 'fragment' and 'light'");
+				return false;
+			}
+			break;
+		default:
+			break;
+	}
+	return true;
+}
+
 bool ShaderLanguage::_validate_assign(Node *p_node, const FunctionInfo &p_function_info, String *r_message) {
 	if (p_node->type == Node::TYPE_OPERATOR) {
 		OperatorNode *op = static_cast<OperatorNode *>(p_node);
@@ -3142,13 +3208,6 @@ bool ShaderLanguage::_validate_assign(Node *p_node, const FunctionInfo &p_functi
 			return false;
 		}
 
-		if (shader->varyings.has(var->name) && current_function != String("vertex")) {
-			if (r_message) {
-				*r_message = RTR("Varyings can only be assigned in vertex function.");
-			}
-			return false;
-		}
-
 		if (shader->constants.has(var->name) || var->is_const) {
 			if (r_message) {
 				*r_message = RTR("Constants cannot be modified.");
@@ -3169,13 +3228,6 @@ bool ShaderLanguage::_validate_assign(Node *p_node, const FunctionInfo &p_functi
 			return false;
 		}
 
-		if (shader->varyings.has(arr->name) && current_function != String("vertex")) {
-			if (r_message) {
-				*r_message = RTR("Varyings can only be assigned in vertex function.");
-			}
-			return false;
-		}
-
 		return true;
 	}
 
@@ -3761,6 +3813,23 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 						_set_error("Unknown identifier in expression: " + String(identifier));
 						return nullptr;
 					}
+					if (ident_type == IDENTIFIER_VARYING) {
+						TkPos prev_pos = _get_tkpos();
+						Token next_token = _get_token();
+						_set_tkpos(prev_pos);
+						String error;
+						if (next_token.type == TK_OP_ASSIGN) {
+							if (!_validate_varying_assign(shader->varyings[identifier], &error)) {
+								_set_error(error);
+								return nullptr;
+							}
+						} else {
+							if (!_validate_varying_using(shader->varyings[identifier], &error)) {
+								_set_error(error);
+								return nullptr;
+							}
+						}
+					}
 					last_const = is_const;
 
 					if (ident_type == IDENTIFIER_FUNCTION) {
@@ -3786,10 +3855,6 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 							_set_error("Constants cannot be modified.");
 							return nullptr;
 						}
-						if (shader->varyings.has(identifier) && current_function != String("vertex")) {
-							_set_error("Varyings can only be assigned in vertex function.");
-							return nullptr;
-						}
 						assign_expression = _parse_array_constructor(p_block, p_function_info, data_type, struct_name, array_size);
 						if (!assign_expression) {
 							return nullptr;
@@ -6260,6 +6325,7 @@ Error ShaderLanguage::_parse_shader(const Map<StringName, FunctionInfo> &p_funct
 					return ERR_PARSE_ERROR;
 				}
 
+				TkPos name_pos = _get_tkpos();
 				name = tk.text;
 
 				if (_find_identifier(nullptr, false, FunctionInfo(), name)) {
@@ -6541,11 +6607,12 @@ Error ShaderLanguage::_parse_shader(const Map<StringName, FunctionInfo> &p_funct
 						_set_error("Expected ';'");
 						return ERR_PARSE_ERROR;
 					}
-				} else {
+				} else { // varying
 					ShaderNode::Varying varying;
 					varying.type = type;
 					varying.precision = precision;
 					varying.interpolation = interpolation;
+					varying.tkpos = name_pos;
 
 					tk = _get_token();
 					if (tk.type != TK_SEMICOLON && tk.type != TK_BRACKET_OPEN) {
@@ -7158,6 +7225,14 @@ Error ShaderLanguage::_parse_shader(const Map<StringName, FunctionInfo> &p_funct
 		tk = _get_token();
 	}
 
+	for (Map<StringName, ShaderNode::Varying>::Element *E = shader->varyings.front(); E; E = E->next()) {
+		if (E->get().stage == ShaderNode::Varying::STAGE_VERTEX || E->get().stage == ShaderNode::Varying::STAGE_FRAGMENT) {
+			_set_tkpos(E->get().tkpos);
+			_set_error(RTR("Varying must only be used in two different stages, which can be 'vertex' 'fragment' and 'light'"));
+			return ERR_PARSE_ERROR;
+		}
+	}
+
 	return OK;
 }
 
diff --git a/servers/rendering/shader_language.h b/servers/rendering/shader_language.h
index 27767378f9..14594b039c 100644
--- a/servers/rendering/shader_language.h
+++ b/servers/rendering/shader_language.h
@@ -41,6 +41,11 @@
 
 class ShaderLanguage {
 public:
+	struct TkPos {
+		int char_idx;
+		int tk_line;
+	};
+
 	enum TokenType {
 		TK_EMPTY,
 		TK_IDENTIFIER,
@@ -598,10 +603,21 @@ public:
 		};
 
 		struct Varying {
+			enum Stage {
+				STAGE_UNKNOWN,
+				STAGE_VERTEX, // transition stage to STAGE_VERTEX_TO_FRAGMENT or STAGE_VERTEX_TO_LIGHT, emits error if they are not used
+				STAGE_FRAGMENT, // transition stage to STAGE_FRAGMENT_TO_LIGHT, emits error if it's not used
+				STAGE_VERTEX_TO_FRAGMENT,
+				STAGE_VERTEX_TO_LIGHT,
+				STAGE_FRAGMENT_TO_LIGHT,
+			};
+
+			Stage stage = STAGE_UNKNOWN;
 			DataType type = TYPE_VOID;
 			DataInterpolation interpolation = INTERPOLATION_FLAT;
 			DataPrecision precision = PRECISION_DEFAULT;
 			int array_size = 0;
+			TkPos tkpos;
 
 			Varying() {}
 		};
@@ -780,11 +796,6 @@ private:
 	StringName current_function;
 	bool last_const = false;
 
-	struct TkPos {
-		int char_idx;
-		int tk_line;
-	};
-
 	TkPos _get_tkpos() {
 		TkPos tkp;
 		tkp.char_idx = char_idx;
@@ -864,6 +875,8 @@ private:
 	bool _parse_function_arguments(BlockNode *p_block, const FunctionInfo &p_function_info, OperatorNode *p_func, int *r_complete_arg = nullptr);
 	bool _propagate_function_call_sampler_uniform_settings(StringName p_name, int p_argument, TextureFilter p_filter, TextureRepeat p_repeat);
 	bool _propagate_function_call_sampler_builtin_reference(StringName p_name, int p_argument, const StringName &p_builtin);
+	bool _validate_varying_assign(ShaderNode::Varying &p_varying, String *r_message);
+	bool _validate_varying_using(ShaderNode::Varying &p_varying, String *r_message);
 
 	Node *_parse_expression(BlockNode *p_block, const FunctionInfo &p_function_info);
 	Node *_parse_array_constructor(BlockNode *p_block, const FunctionInfo &p_function_info, DataType p_type, const StringName &p_struct_name, int p_array_size);