Implement coloring for disabled branches in the shader editor

2 files changed, 66 insertions, 12 deletions

diff --git a/servers/rendering/shader_preprocessor.cpp b/servers/rendering/shader_preprocessor.cpp
index a7b274b3e2..d118c73c4a 100644
--- a/servers/rendering/shader_preprocessor.cpp
+++ b/servers/rendering/shader_preprocessor.cpp
@@ -416,16 +416,22 @@ void ShaderPreprocessor::process_define(Tokenizer *p_tokenizer) {
 }
 
 void ShaderPreprocessor::process_else(Tokenizer *p_tokenizer) {
+	const int line = p_tokenizer->get_line();
+
 	if (state->skip_stack_else.is_empty()) {
-		set_error(RTR("Unmatched else."), p_tokenizer->get_line());
+		set_error(RTR("Unmatched else."), line);
 		return;
 	}
+	if (state->previous_region != nullptr) {
+		state->previous_region->to_line = line - 1;
+	}
+
 	p_tokenizer->advance('\n');
 
 	bool skip = state->skip_stack_else[state->skip_stack_else.size() - 1];
 	state->skip_stack_else.remove_at(state->skip_stack_else.size() - 1);
 
-	Vector<SkippedCondition *> vec = state->skipped_conditions[state->current_include];
+	Vector<SkippedCondition *> vec = state->skipped_conditions[state->current_filename];
 	int index = vec.size() - 1;
 	if (index >= 0) {
 		SkippedCondition *cond = vec[index];
@@ -434,6 +440,10 @@ void ShaderPreprocessor::process_else(Tokenizer *p_tokenizer) {
 		}
 	}
 
+	if (state->save_regions) {
+		add_region(line + 1, !skip, state->previous_region->parent);
+	}
+
 	if (skip) {
 		Vector<String> ends;
 		ends.push_back("endif");
@@ -447,8 +457,12 @@ void ShaderPreprocessor::process_endif(Tokenizer *p_tokenizer) {
 		set_error(RTR("Unmatched endif."), p_tokenizer->get_line());
 		return;
 	}
+	if (state->previous_region != nullptr) {
+		state->previous_region->to_line = p_tokenizer->get_line() - 1;
+		state->previous_region = state->previous_region->parent;
+	}
 
-	Vector<SkippedCondition *> vec = state->skipped_conditions[state->current_include];
+	Vector<SkippedCondition *> vec = state->skipped_conditions[state->current_filename];
 	int index = vec.size() - 1;
 	if (index >= 0) {
 		SkippedCondition *cond = vec[index];
@@ -461,7 +475,7 @@ void ShaderPreprocessor::process_endif(Tokenizer *p_tokenizer) {
 }
 
 void ShaderPreprocessor::process_if(Tokenizer *p_tokenizer) {
-	int line = p_tokenizer->get_line();
+	const int line = p_tokenizer->get_line();
 
 	String body = tokens_to_string(p_tokenizer->advance('\n')).strip_edges();
 	if (body.is_empty()) {
@@ -490,6 +504,10 @@ void ShaderPreprocessor::process_if(Tokenizer *p_tokenizer) {
 
 	bool success = v.booleanize();
 	start_branch_condition(p_tokenizer, success);
+
+	if (state->save_regions) {
+		add_region(line + 1, success, state->previous_region);
+	}
 }
 
 void ShaderPreprocessor::process_ifdef(Tokenizer *p_tokenizer) {
@@ -510,6 +528,10 @@ void ShaderPreprocessor::process_ifdef(Tokenizer *p_tokenizer) {
 
 	bool success = state->defines.has(label);
 	start_branch_condition(p_tokenizer, success);
+
+	if (state->save_regions) {
+		add_region(line + 1, success, state->previous_region);
+	}
 }
 
 void ShaderPreprocessor::process_ifndef(Tokenizer *p_tokenizer) {
@@ -530,6 +552,10 @@ void ShaderPreprocessor::process_ifndef(Tokenizer *p_tokenizer) {
 
 	bool success = !state->defines.has(label);
 	start_branch_condition(p_tokenizer, success);
+
+	if (state->save_regions) {
+		add_region(line + 1, success, state->previous_region);
+	}
 }
 
 void ShaderPreprocessor::process_include(Tokenizer *p_tokenizer) {
@@ -594,15 +620,15 @@ void ShaderPreprocessor::process_include(Tokenizer *p_tokenizer) {
 		return;
 	}
 
-	String old_include = state->current_include;
-	state->current_include = real_path;
+	String old_filename = state->current_filename;
+	state->current_filename = real_path;
 	ShaderPreprocessor processor;
 
 	int prev_condition_depth = state->condition_depth;
 	state->condition_depth = 0;
 
 	FilePosition fp;
-	fp.file = state->current_include;
+	fp.file = state->current_filename;
 	fp.line = line;
 	state->include_positions.push_back(fp);
 
@@ -614,7 +640,7 @@ void ShaderPreprocessor::process_include(Tokenizer *p_tokenizer) {
 
 	// Reset to last include if there are no errors. We want to use this as context.
 	if (state->error.is_empty()) {
-		state->current_include = old_include;
+		state->current_filename = old_filename;
 		state->include_positions.pop_back();
 	} else {
 		return;
@@ -668,6 +694,15 @@ void ShaderPreprocessor::process_undef(Tokenizer *p_tokenizer) {
 	state->defines.erase(label);
 }
 
+void ShaderPreprocessor::add_region(int p_line, bool p_enabled, Region *p_parent_region) {
+	Region region;
+	region.file = state->current_filename;
+	region.enabled = p_enabled;
+	region.from_line = p_line;
+	region.parent = p_parent_region;
+	state->previous_region = &state->regions[region.file].push_back(region)->get();
+}
+
 void ShaderPreprocessor::start_branch_condition(Tokenizer *p_tokenizer, bool p_success) {
 	state->condition_depth++;
 
@@ -676,7 +711,7 @@ void ShaderPreprocessor::start_branch_condition(Tokenizer *p_tokenizer, bool p_s
 	} else {
 		SkippedCondition *cond = memnew(SkippedCondition());
 		cond->start_line = p_tokenizer->get_line();
-		state->skipped_conditions[state->current_include].push_back(cond);
+		state->skipped_conditions[state->current_filename].push_back(cond);
 
 		Vector<String> ends;
 		ends.push_back("else");
@@ -969,8 +1004,12 @@ Error ShaderPreprocessor::preprocess(State *p_state, const String &p_code, Strin
 	return OK;
 }
 
-Error ShaderPreprocessor::preprocess(const String &p_code, String &r_result, String *r_error_text, List<FilePosition> *r_error_position, HashSet<Ref<ShaderInclude>> *r_includes, List<ScriptLanguage::CodeCompletionOption> *r_completion_options, IncludeCompletionFunction p_include_completion_func) {
+Error ShaderPreprocessor::preprocess(const String &p_code, const String &p_filename, String &r_result, String *r_error_text, List<FilePosition> *r_error_position, List<Region> *r_regions, HashSet<Ref<ShaderInclude>> *r_includes, List<ScriptLanguage::CodeCompletionOption> *r_completion_options, IncludeCompletionFunction p_include_completion_func) {
 	State pp_state;
+	if (!p_filename.is_empty()) {
+		pp_state.current_filename = p_filename;
+		pp_state.save_regions = r_regions != nullptr;
+	}
 	Error err = preprocess(&pp_state, p_code, r_result);
 	if (err != OK) {
 		if (r_error_text) {
@@ -980,6 +1019,9 @@ Error ShaderPreprocessor::preprocess(const String &p_code, String &r_result, Str
 			*r_error_position = pp_state.include_positions;
 		}
 	}
+	if (r_regions) {
+		*r_regions = pp_state.regions[p_filename];
+	}
 	if (r_includes) {
 		*r_includes = pp_state.shader_includes;
 	}
diff --git a/servers/rendering/shader_preprocessor.h b/servers/rendering/shader_preprocessor.h
index a93fb680dd..41b574298d 100644
--- a/servers/rendering/shader_preprocessor.h
+++ b/servers/rendering/shader_preprocessor.h
@@ -58,6 +58,14 @@ public:
 		int line = 0;
 	};
 
+	struct Region {
+		String file;
+		int from_line = -1;
+		int to_line = -1;
+		bool enabled = false;
+		Region *parent = nullptr;
+	};
+
 private:
 	struct Token {
 		char32_t text;
@@ -134,10 +142,13 @@ private:
 		RBSet<String> includes;
 		List<uint64_t> cyclic_include_hashes; // Holds code hash of includes.
 		int include_depth = 0;
-		String current_include;
+		String current_filename;
 		String current_shader_type;
 		String error;
 		List<FilePosition> include_positions;
+		bool save_regions = false;
+		RBMap<String, List<Region>> regions;
+		Region *previous_region = nullptr;
 		RBMap<String, Vector<SkippedCondition *>> skipped_conditions;
 		bool disabled = false;
 		CompletionType completion_type = COMPLETION_TYPE_NONE;
@@ -167,6 +178,7 @@ private:
 	void process_pragma(Tokenizer *p_tokenizer);
 	void process_undef(Tokenizer *p_tokenizer);
 
+	void add_region(int p_line, bool p_enabled, Region *p_parent_region);
 	void start_branch_condition(Tokenizer *p_tokenizer, bool p_success);
 
 	void expand_output_macros(int p_start, int p_line);
@@ -188,7 +200,7 @@ private:
 public:
 	typedef void (*IncludeCompletionFunction)(List<ScriptLanguage::CodeCompletionOption> *);
 
-	Error preprocess(const String &p_code, String &r_result, String *r_error_text = nullptr, List<FilePosition> *r_error_position = nullptr, HashSet<Ref<ShaderInclude>> *r_includes = nullptr, List<ScriptLanguage::CodeCompletionOption> *r_completion_options = nullptr, IncludeCompletionFunction p_include_completion_func = nullptr);
+	Error preprocess(const String &p_code, const String &p_filename, String &r_result, String *r_error_text = nullptr, List<FilePosition> *r_error_position = nullptr, List<Region> *r_regions = nullptr, HashSet<Ref<ShaderInclude>> *r_includes = nullptr, List<ScriptLanguage::CodeCompletionOption> *r_completion_options = nullptr, IncludeCompletionFunction p_include_completion_func = nullptr);
 
 	static void get_keyword_list(List<String> *r_keywords, bool p_include_shader_keywords);
 	static void get_pragma_list(List<String> *r_pragmas);