Fix shader language preprocessor include marker handling

3 files changed, 33 insertions, 10 deletions

diff --git a/servers/rendering/shader_language.cpp b/servers/rendering/shader_language.cpp
index 1460a91826..d6d6cc7731 100644
--- a/servers/rendering/shader_language.cpp
+++ b/servers/rendering/shader_language.cpp
@@ -610,19 +610,39 @@ ShaderLanguage::Token ShaderLanguage::_get_token() {
 					incp.push_back(0); // Zero end it.
 					String include_path(incp.ptr());
 					include_positions.write[include_positions.size() - 1].line = tk_line;
-					FilePosition fp;
-					fp.file = include_path;
-					fp.line = 0;
-					tk_line = 0;
-					include_positions.push_back(fp);
 
-				} else if (GETCHAR(0) == '@' && GETCHAR(1) == '<') {
-					if (include_positions.size() == 1) {
-						return _make_token(TK_ERROR, "Invalid include exit hint @@< without matching enter hint.");
+					String marker = ">>" + include_path;
+					if (!include_markers_handled.has(marker)) {
+						include_markers_handled.insert(marker);
+
+						FilePosition fp;
+						fp.file = include_path;
+						fp.line = 0;
+						tk_line = 0;
+						include_positions.push_back(fp);
 					}
+
+				} else if (GETCHAR(0) == '@' && GETCHAR(1) == '<') {
 					char_idx += 2;
 
-					include_positions.resize(include_positions.size() - 1); // Pop back.
+					LocalVector<char32_t> incp;
+					while (GETCHAR(0) != '\n') {
+						incp.push_back(GETCHAR(0));
+						char_idx++;
+					}
+					incp.push_back(0); // Zero end it.
+					String include_path(incp.ptr());
+
+					String marker = "<<" + include_path;
+					if (!include_markers_handled.has(marker)) {
+						include_markers_handled.insert(marker);
+
+						if (include_positions.size() == 1) {
+							return _make_token(TK_ERROR, "Invalid include exit hint @@< without matching enter hint.");
+						}
+						include_positions.resize(include_positions.size() - 1); // Pop back.
+					}
+
 					tk_line = include_positions[include_positions.size() - 1].line - 1; // Restore line.
 
 				} else {
@@ -1217,6 +1237,8 @@ void ShaderLanguage::clear() {
 	include_positions.clear();
 	include_positions.push_back(FilePosition());
 
+	include_markers_handled.clear();
+
 #ifdef DEBUG_ENABLED
 	keyword_completion_context = CF_UNSPECIFIED;
 	used_constants.clear();
diff --git a/servers/rendering/shader_language.h b/servers/rendering/shader_language.h
index 0ddd27f028..90fde72152 100644
--- a/servers/rendering/shader_language.h
+++ b/servers/rendering/shader_language.h
@@ -910,6 +910,7 @@ private:
 	int error_line = 0;
 
 	Vector<FilePosition> include_positions;
+	HashSet<String> include_markers_handled;
 
 #ifdef DEBUG_ENABLED
 	struct Usage {
diff --git a/servers/rendering/shader_preprocessor.cpp b/servers/rendering/shader_preprocessor.cpp
index 8d6ddf9121..dbd1374941 100644
--- a/servers/rendering/shader_preprocessor.cpp
+++ b/servers/rendering/shader_preprocessor.cpp
@@ -747,7 +747,7 @@ void ShaderPreprocessor::process_include(Tokenizer *p_tokenizer) {
 	processor.preprocess(state, included, result);
 	add_to_output("@@>" + real_path + "\n"); // Add token for enter include path
 	add_to_output(result);
-	add_to_output("\n@@<\n"); // Add token for exit include path
+	add_to_output("\n@@<" + real_path + "\n"); // Add token for exit include path.
 
 	// Reset to last include if there are no errors. We want to use this as context.
 	if (state->error.is_empty()) {