Fix shader crash when using a varying in separate func before it defined

2 files changed, 34 insertions, 9 deletions

diff --git a/servers/rendering/shader_language.cpp b/servers/rendering/shader_language.cpp
index e35fde406f..745dcf5392 100644
--- a/servers/rendering/shader_language.cpp
+++ b/servers/rendering/shader_language.cpp
@@ -4664,9 +4664,18 @@ bool ShaderLanguage::_validate_restricted_func(const StringName &p_name, const C
 		}
 	}
 
-	if (!p_func_info->uses_restricted_functions.is_empty()) {
-		const Pair<StringName, TkPos> &first_element = p_func_info->uses_restricted_functions.get(0);
-		_set_tkpos(first_element.second);
+	if (!p_func_info->uses_restricted_items.is_empty()) {
+		const Pair<StringName, CallInfo::Item> &first_element = p_func_info->uses_restricted_items.get(0);
+
+		if (first_element.second.type == CallInfo::Item::ITEM_TYPE_VARYING) {
+			const ShaderNode::Varying &varying = shader->varyings[first_element.first];
+
+			if (varying.stage == ShaderNode::Varying::STAGE_VERTEX) {
+				return true;
+			}
+		}
+
+		_set_tkpos(first_element.second.pos);
 
 		if (is_in_restricted_function) {
 			_set_error(vformat(RTR("'%s' cannot be used within the '%s' processor function."), first_element.first, "vertex"));
@@ -5352,7 +5361,7 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 									return nullptr;
 								}
 								// Register usage of the restricted function.
-								calls_info[current_function].uses_restricted_functions.push_back(Pair<StringName, TkPos>(name, _get_tkpos()));
+								calls_info[current_function].uses_restricted_items.push_back(Pair<StringName, CallInfo::Item>(name, CallInfo::Item(CallInfo::Item::ITEM_TYPE_BUILTIN, _get_tkpos())));
 								is_builtin = true;
 								break;
 							}
@@ -5469,10 +5478,11 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 
 										if (shader->varyings.has(varname)) {
 											switch (shader->varyings[varname].stage) {
-												case ShaderNode::Varying::STAGE_UNKNOWN: {
-													_set_error(vformat(RTR("Varying '%s' must be assigned in the 'vertex' or 'fragment' function first."), varname));
-													return nullptr;
-												}
+												case ShaderNode::Varying::STAGE_UNKNOWN:
+													if (is_out_arg) {
+														error = true;
+													}
+													break;
 												case ShaderNode::Varying::STAGE_VERTEX_TO_FRAGMENT_LIGHT:
 													[[fallthrough]];
 												case ShaderNode::Varying::STAGE_VERTEX:
@@ -5672,6 +5682,8 @@ ShaderLanguage::Node *ShaderLanguage::_parse_expression(BlockNode *p_block, cons
 						_set_tkpos(prev_pos);
 
 						ShaderNode::Varying &var = shader->varyings[identifier];
+						calls_info[current_function].uses_restricted_items.push_back(Pair<StringName, CallInfo::Item>(identifier, CallInfo::Item(CallInfo::Item::ITEM_TYPE_VARYING, prev_pos)));
+
 						String error;
 						if (is_token_operator_assign(next_token.type)) {
 							if (!_validate_varying_assign(shader->varyings[identifier], &error)) {
diff --git a/servers/rendering/shader_language.h b/servers/rendering/shader_language.h
index 076bd8def4..bc0aa0558a 100644
--- a/servers/rendering/shader_language.h
+++ b/servers/rendering/shader_language.h
@@ -917,8 +917,21 @@ private:
 
 	// Additional function information (eg. call hierarchy). No need to expose it to compiler.
 	struct CallInfo {
+		struct Item {
+			enum ItemType {
+				ITEM_TYPE_BUILTIN,
+				ITEM_TYPE_VARYING,
+			} type;
+
+			TkPos pos;
+
+			Item() {}
+			Item(ItemType p_type, TkPos p_pos) :
+					type(p_type), pos(p_pos) {}
+		};
+
 		StringName name;
-		List<Pair<StringName, TkPos>> uses_restricted_functions;
+		List<Pair<StringName, Item>> uses_restricted_items;
 		List<CallInfo *> calls;
 	};