Several shader preprocessor parser fixes and improvements

3 files changed, 114 insertions, 35 deletions

diff --git a/servers/rendering/shader_language.cpp b/servers/rendering/shader_language.cpp
index 91e9879bae..131b16757e 100644
--- a/servers/rendering/shader_language.cpp
+++ b/servers/rendering/shader_language.cpp
@@ -622,7 +622,7 @@ ShaderLanguage::Token ShaderLanguage::_get_token() {
 					char_idx += 2;
 
 					include_positions.resize(include_positions.size() - 1); // Pop back.
-					tk_line = include_positions[include_positions.size() - 1].line; // Restore line.
+					tk_line = include_positions[include_positions.size() - 1].line - 1; // Restore line.
 
 				} else {
 					return _make_token(TK_ERROR, "Invalid include enter/exit hint token (@@> and @@<)");
diff --git a/servers/rendering/shader_preprocessor.cpp b/servers/rendering/shader_preprocessor.cpp
index 97efd2df94..816d937e9f 100644
--- a/servers/rendering/shader_preprocessor.cpp
+++ b/servers/rendering/shader_preprocessor.cpp
@@ -78,19 +78,46 @@ char32_t ShaderPreprocessor::Tokenizer::peek() {
 	return 0;
 }
 
+int ShaderPreprocessor::Tokenizer::consume_line_continuations(int p_offset) {
+	int skips = 0;
+
+	for (int i = index + p_offset; i < size; i++) {
+		char32_t c = code[i];
+		if (c == '\\') {
+			if (i + 1 < size && code[i + 1] == '\n') {
+				// This line ends with "\" and "\n" continuation.
+				add_generated(Token('\n', line));
+				line++;
+				skips++;
+
+				i = i + 2;
+				index = i;
+			} else {
+				break;
+			}
+		} else if (!is_whitespace(c)) {
+			break;
+		}
+	}
+	return skips;
+}
+
 LocalVector<ShaderPreprocessor::Token> ShaderPreprocessor::Tokenizer::advance(char32_t p_what) {
 	LocalVector<ShaderPreprocessor::Token> tokens;
 
 	while (index < size) {
 		char32_t c = code[index++];
-
-		tokens.push_back(ShaderPreprocessor::Token(c, line));
+		if (c == '\\' && consume_line_continuations(-1) > 0) {
+			continue;
+		}
 
 		if (c == '\n') {
 			add_generated(ShaderPreprocessor::Token('\n', line));
 			line++;
 		}
 
+		tokens.push_back(ShaderPreprocessor::Token(c, line));
+
 		if (c == p_what || c == 0) {
 			return tokens;
 		}
@@ -104,6 +131,11 @@ void ShaderPreprocessor::Tokenizer::skip_whitespace() {
 	}
 }
 
+bool ShaderPreprocessor::Tokenizer::consume_empty_line() {
+	// Read until newline and return true if the content was all whitespace/empty.
+	return tokens_to_string(advance('\n')).strip_edges().size() == 0;
+}
+
 String ShaderPreprocessor::Tokenizer::get_identifier(bool *r_is_cursor, bool p_started) {
 	if (r_is_cursor != nullptr) {
 		*r_is_cursor = false;
@@ -113,6 +145,10 @@ String ShaderPreprocessor::Tokenizer::get_identifier(bool *r_is_cursor, bool p_s
 
 	while (true) {
 		char32_t c = peek();
+		if (c == '\\' && consume_line_continuations(0) > 0) {
+			continue;
+		}
+
 		if (is_char_end(c) || c == '(' || c == ')' || c == ',' || c == ';') {
 			break;
 		}
@@ -146,8 +182,10 @@ String ShaderPreprocessor::Tokenizer::get_identifier(bool *r_is_cursor, bool p_s
 
 String ShaderPreprocessor::Tokenizer::peek_identifier() {
 	const int original = index;
+	const int original_line = line;
 	String id = get_identifier();
 	index = original;
+	line = original_line;
 	return id;
 }
 
@@ -485,7 +523,9 @@ void ShaderPreprocessor::process_else(Tokenizer *p_tokenizer) {
 		state->previous_region->to_line = line - 1;
 	}
 
-	p_tokenizer->advance('\n');
+	if (!p_tokenizer->consume_empty_line()) {
+		set_error(RTR("Invalid else."), p_tokenizer->get_line());
+	}
 
 	bool skip = false;
 	for (int i = 0; i < state->current_branch->conditions.size(); i++) {
@@ -508,17 +548,21 @@ void ShaderPreprocessor::process_else(Tokenizer *p_tokenizer) {
 }
 
 void ShaderPreprocessor::process_endif(Tokenizer *p_tokenizer) {
+	const int line = p_tokenizer->get_line();
+
 	state->condition_depth--;
 	if (state->condition_depth < 0) {
-		set_error(RTR("Unmatched endif."), p_tokenizer->get_line());
+		set_error(RTR("Unmatched endif."), line);
 		return;
 	}
 	if (state->previous_region != nullptr) {
-		state->previous_region->to_line = p_tokenizer->get_line() - 1;
+		state->previous_region->to_line = line - 1;
 		state->previous_region = state->previous_region->parent;
 	}
 
-	p_tokenizer->advance('\n');
+	if (!p_tokenizer->consume_empty_line()) {
+		set_error(RTR("Invalid endif."), line);
+	}
 
 	state->current_branch = state->current_branch->parent;
 	state->branches.pop_back();
@@ -574,12 +618,10 @@ void ShaderPreprocessor::process_ifdef(Tokenizer *p_tokenizer) {
 		return;
 	}
 
-	p_tokenizer->skip_whitespace();
-	if (!is_char_end(p_tokenizer->peek())) {
+	if (!p_tokenizer->consume_empty_line()) {
 		set_error(RTR("Invalid ifdef."), line);
 		return;
 	}
-	p_tokenizer->advance('\n');
 
 	bool success = state->defines.has(label);
 	start_branch_condition(p_tokenizer, success);
@@ -598,12 +640,10 @@ void ShaderPreprocessor::process_ifndef(Tokenizer *p_tokenizer) {
 		return;
 	}
 
-	p_tokenizer->skip_whitespace();
-	if (!is_char_end(p_tokenizer->peek())) {
+	if (!p_tokenizer->consume_empty_line()) {
 		set_error(RTR("Invalid ifndef."), line);
 		return;
 	}
-	p_tokenizer->advance('\n');
 
 	bool success = !state->defines.has(label);
 	start_branch_condition(p_tokenizer, success);
@@ -628,9 +668,8 @@ void ShaderPreprocessor::process_include(Tokenizer *p_tokenizer) {
 		}
 	}
 	path = path.substr(0, path.length() - 1);
-	p_tokenizer->skip_whitespace();
 
-	if (path.is_empty() || !is_char_end(p_tokenizer->peek())) {
+	if (path.is_empty() || !p_tokenizer->consume_empty_line()) {
 		set_error(RTR("Invalid path."), line);
 		return;
 	}
@@ -728,25 +767,24 @@ void ShaderPreprocessor::process_pragma(Tokenizer *p_tokenizer) {
 		return;
 	}
 
-	p_tokenizer->advance('\n');
+	if (!p_tokenizer->consume_empty_line()) {
+		set_error(RTR("Invalid pragma directive."), line);
+		return;
+	}
 }
 
 void ShaderPreprocessor::process_undef(Tokenizer *p_tokenizer) {
 	const int line = p_tokenizer->get_line();
 	const String label = p_tokenizer->get_identifier();
-	if (label.is_empty() || !state->defines.has(label)) {
-		set_error(RTR("Invalid name."), line);
-		return;
-	}
-
-	p_tokenizer->skip_whitespace();
-	if (!is_char_end(p_tokenizer->peek())) {
+	if (label.is_empty() || !p_tokenizer->consume_empty_line()) {
 		set_error(RTR("Invalid undef."), line);
 		return;
 	}
 
-	memdelete(state->defines[label]);
-	state->defines.erase(label);
+	if (state->defines.has(label)) {
+		memdelete(state->defines[label]);
+		state->defines.erase(label);
+	}
 }
 
 void ShaderPreprocessor::add_region(int p_line, bool p_enabled, Region *p_parent_region) {
@@ -957,15 +995,57 @@ bool ShaderPreprocessor::expand_macros_once(const String &p_line, int p_line_num
 		String body = define->body;
 		if (define->arguments.size() > 0) {
 			// Complex macro with arguments.
-			int args_start = index + key.length();
-			int args_end = p_line.find(")", args_start);
-			if (args_start == -1 || args_end == -1) {
-				set_error(RTR("Missing macro argument parenthesis."), p_line_number);
-				return false;
+
+			int args_start = -1;
+			int args_end = -1;
+			int brackets_open = 0;
+			Vector<String> args;
+			for (int i = index_start - 1; i < p_line.length(); i++) {
+				bool add_argument = false;
+				bool reached_end = false;
+				char32_t c = p_line[i];
+
+				if (c == '(') {
+					brackets_open++;
+					if (brackets_open == 1) {
+						args_start = i + 1;
+						args_end = -1;
+					}
+				} else if (c == ')') {
+					brackets_open--;
+					if (brackets_open == 0) {
+						args_end = i;
+						add_argument = true;
+						reached_end = true;
+					}
+				} else if (c == ',') {
+					if (brackets_open == 1) {
+						args_end = i;
+						add_argument = true;
+					}
+				}
+
+				if (add_argument) {
+					if (args_start == -1 || args_end == -1) {
+						set_error(RTR("Invalid macro argument list."), p_line_number);
+						return false;
+					}
+
+					String arg = p_line.substr(args_start, args_end - args_start).strip_edges();
+					if (arg.is_empty()) {
+						set_error(RTR("Invalid macro argument."), p_line_number);
+						return false;
+					}
+					args.append(arg);
+
+					args_start = args_end + 1;
+				}
+
+				if (reached_end) {
+					break;
+				}
 			}
 
-			String values = result.substr(args_start + 1, args_end - (args_start + 1));
-			Vector<String> args = values.split(",");
 			if (args.size() != define->arguments.size()) {
 				set_error(RTR("Invalid macro argument count."), p_line_number);
 				return false;
@@ -987,9 +1067,6 @@ bool ShaderPreprocessor::expand_macros_once(const String &p_line, int p_line_num
 			result = result.substr(0, index) + " " + body + " " + result.substr(args_end + 1, result.length());
 		} else {
 			result = result.substr(0, index) + body + result.substr(index + key.length(), result.length() - (index + key.length()));
-			// Manually reset index_start to where the body value of the define finishes.
-			// This ensures we don't skip another instance of this macro in the string.
-			index_start = index + body.length() + 1;
 		}
 
 		r_expanded = result;
diff --git a/servers/rendering/shader_preprocessor.h b/servers/rendering/shader_preprocessor.h
index 6e5533c575..2b7d2c274e 100644
--- a/servers/rendering/shader_preprocessor.h
+++ b/servers/rendering/shader_preprocessor.h
@@ -93,11 +93,13 @@ private:
 		int get_line() const;
 		int get_index() const;
 		char32_t peek();
+		int consume_line_continuations(int p_offset);
 
 		void get_and_clear_generated(Vector<Token> *r_out);
 		void backtrack(char32_t p_what);
 		LocalVector<Token> advance(char32_t p_what);
 		void skip_whitespace();
+		bool consume_empty_line();
 		String get_identifier(bool *r_is_cursor = nullptr, bool p_started = false);
 		String peek_identifier();
 		Token get_token();