1 files changed, 66 insertions, 48 deletions

diff --git a/servers/rendering/renderer_rd/shader_rd.cpp b/servers/rendering/renderer_rd/shader_rd.cpp
index 39c3e9b168..6234cddee3 100644
--- a/servers/rendering/renderer_rd/shader_rd.cpp
+++ b/servers/rendering/renderer_rd/shader_rd.cpp
@@ -39,6 +39,8 @@
 #include "servers/rendering/rendering_device.h"
 #include "thirdparty/misc/smolv.h"
 
+#define ENABLE_SHADER_CACHE 1
+
 void ShaderRD::_add_stage(const char *p_code, StageType p_stage_type) {
 	Vector<String> lines = String(p_code).split("\n");
 
@@ -144,7 +146,8 @@ RID ShaderRD::version_create() {
 	version.dirty = true;
 	version.valid = false;
 	version.initialize_needed = true;
-	version.variants = nullptr;
+	version.variants.clear();
+	version.variant_data.clear();
 	return version_owner.make_rid(version);
 }
 
@@ -154,23 +157,25 @@ void ShaderRD::_initialize_version(Version *p_version) {
 	p_version->valid = false;
 	p_version->dirty = false;
 
-	p_version->variants = memnew_arr(RID, variant_defines.size());
+	p_version->variants.resize_zeroed(variant_defines.size());
+	p_version->variant_data.resize(variant_defines.size());
+	p_version->group_compilation_tasks.resize(group_enabled.size());
+	p_version->group_compilation_tasks.fill(0);
 }
 
 void ShaderRD::_clear_version(Version *p_version) {
+	_compile_ensure_finished(p_version);
+
 	// Clear versions if they exist.
-	if (p_version->variants) {
+	if (!p_version->variants.is_empty()) {
 		for (int i = 0; i < variant_defines.size(); i++) {
 			if (p_version->variants[i].is_valid()) {
 				RD::get_singleton()->free(p_version->variants[i]);
 			}
 		}
 
-		memdelete_arr(p_version->variants);
-		if (p_version->variant_data) {
-			memdelete_arr(p_version->variant_data);
-		}
-		p_version->variants = nullptr;
+		p_version->variants.clear();
+		p_version->variant_data.clear();
 	}
 }
 
@@ -227,8 +232,8 @@ void ShaderRD::_build_variant_code(StringBuilder &builder, uint32_t p_variant, c
 	}
 }
 
-void ShaderRD::_compile_variant(uint32_t p_variant, const CompileData *p_data) {
-	uint32_t variant = group_to_variant_map[p_data->group][p_variant];
+void ShaderRD::_compile_variant(uint32_t p_variant, CompileData p_data) {
+	uint32_t variant = group_to_variant_map[p_data.group][p_variant];
 
 	if (!variants_enabled[variant]) {
 		return; // Variant is disabled, return.
@@ -245,7 +250,7 @@ void ShaderRD::_compile_variant(uint32_t p_variant, const CompileData *p_data) {
 		//vertex stage
 
 		StringBuilder builder;
-		_build_variant_code(builder, variant, p_data->version, stage_templates[STAGE_TYPE_VERTEX]);
+		_build_variant_code(builder, variant, p_data.version, stage_templates[STAGE_TYPE_VERTEX]);
 
 		current_source = builder.as_string();
 		RD::ShaderStageSPIRVData stage;
@@ -263,7 +268,7 @@ void ShaderRD::_compile_variant(uint32_t p_variant, const CompileData *p_data) {
 		current_stage = RD::SHADER_STAGE_FRAGMENT;
 
 		StringBuilder builder;
-		_build_variant_code(builder, variant, p_data->version, stage_templates[STAGE_TYPE_FRAGMENT]);
+		_build_variant_code(builder, variant, p_data.version, stage_templates[STAGE_TYPE_FRAGMENT]);
 
 		current_source = builder.as_string();
 		RD::ShaderStageSPIRVData stage;
@@ -281,7 +286,7 @@ void ShaderRD::_compile_variant(uint32_t p_variant, const CompileData *p_data) {
 		current_stage = RD::SHADER_STAGE_COMPUTE;
 
 		StringBuilder builder;
-		_build_variant_code(builder, variant, p_data->version, stage_templates[STAGE_TYPE_COMPUTE]);
+		_build_variant_code(builder, variant, p_data.version, stage_templates[STAGE_TYPE_COMPUTE]);
 
 		current_source = builder.as_string();
 
@@ -313,8 +318,8 @@ void ShaderRD::_compile_variant(uint32_t p_variant, const CompileData *p_data) {
 	{
 		MutexLock lock(variant_set_mutex);
 
-		p_data->version->variants[variant] = RD::get_singleton()->shader_create_from_bytecode(shader_data, p_data->version->variants[variant]);
-		p_data->version->variant_data[variant] = shader_data;
+		p_data.version->variants.write[variant] = RD::get_singleton()->shader_create_from_bytecode(shader_data, p_data.version->variants[variant]);
+		p_data.version->variant_data.write[variant] = shader_data;
 	}
 }
 
@@ -443,14 +448,14 @@ bool ShaderRD::_load_from_cache(Version *p_version, int p_group) {
 
 		ERR_FAIL_COND_V(br != variant_size, false);
 
-		p_version->variant_data[variant_id] = variant_bytes;
+		p_version->variant_data.write[variant_id] = variant_bytes;
 	}
 
 	for (uint32_t i = 0; i < variant_count; i++) {
 		int variant_id = group_to_variant_map[p_group][i];
 		if (!variants_enabled[variant_id]) {
 			MutexLock lock(variant_set_mutex);
-			p_version->variants[variant_id] = RID();
+			p_version->variants.write[variant_id] = RID();
 			continue;
 		}
 		{
@@ -464,12 +469,10 @@ bool ShaderRD::_load_from_cache(Version *p_version, int p_group) {
 				ERR_FAIL_COND_V(shader.is_null(), false);
 			}
 
-			p_version->variants[variant_id] = shader;
+			p_version->variants.write[variant_id] = shader;
 		}
 	}
 
-	memdelete_arr(p_version->variant_data); //clear stages
-	p_version->variant_data = nullptr;
 	p_version->valid = true;
 	return true;
 }
@@ -491,48 +494,51 @@ void ShaderRD::_save_to_cache(Version *p_version, int p_group) {
 }
 
 void ShaderRD::_allocate_placeholders(Version *p_version, int p_group) {
-	ERR_FAIL_NULL(p_version->variants);
+	ERR_FAIL_COND(p_version->variants.is_empty());
+
 	for (uint32_t i = 0; i < group_to_variant_map[p_group].size(); i++) {
 		int variant_id = group_to_variant_map[p_group][i];
 		RID shader = RD::get_singleton()->shader_create_placeholder();
 		{
 			MutexLock lock(variant_set_mutex);
-			p_version->variants[variant_id] = shader;
+			p_version->variants.write[variant_id] = shader;
 		}
 	}
 }
 
 // Try to compile all variants for a given group.
 // Will skip variants that are disabled.
-void ShaderRD::_compile_version(Version *p_version, int p_group) {
+void ShaderRD::_compile_version_start(Version *p_version, int p_group) {
 	if (!group_enabled[p_group]) {
 		return;
 	}
 
-	typedef Vector<uint8_t> ShaderStageData;
-	p_version->variant_data = memnew_arr(ShaderStageData, variant_defines.size());
-
 	p_version->dirty = false;
 
+#if ENABLE_SHADER_CACHE
 	if (shader_cache_dir_valid) {
 		if (_load_from_cache(p_version, p_group)) {
 			return;
 		}
 	}
+#endif
 
 	CompileData compile_data;
 	compile_data.version = p_version;
 	compile_data.group = p_group;
 
-#if 1
-	WorkerThreadPool::GroupID group_task = WorkerThreadPool::get_singleton()->add_template_group_task(this, &ShaderRD::_compile_variant, &compile_data, group_to_variant_map[p_group].size(), -1, true, SNAME("ShaderCompilation"));
-	WorkerThreadPool::get_singleton()->wait_for_group_task_completion(group_task);
+	WorkerThreadPool::GroupID group_task = WorkerThreadPool::get_singleton()->add_template_group_task(this, &ShaderRD::_compile_variant, compile_data, group_to_variant_map[p_group].size(), -1, true, SNAME("ShaderCompilation"));
+	p_version->group_compilation_tasks.write[p_group] = group_task;
+}
 
-#else
-	for (uint32_t i = 0; i < group_to_variant_map[p_group].size(); i++) {
-		_compile_variant(i, &compile_data);
+void ShaderRD::_compile_version_end(Version *p_version, int p_group) {
+	if (p_version->group_compilation_tasks.size() <= p_group || p_version->group_compilation_tasks[p_group] == 0) {
+		return;
 	}
-#endif
+
+	WorkerThreadPool::GroupID group_task = p_version->group_compilation_tasks[p_group];
+	WorkerThreadPool::get_singleton()->wait_for_group_task_completion(group_task);
+	p_version->group_compilation_tasks.write[p_group] = 0;
 
 	bool all_valid = true;
 
@@ -557,29 +563,35 @@ void ShaderRD::_compile_version(Version *p_version, int p_group) {
 				RD::get_singleton()->free(p_version->variants[i]);
 			}
 		}
-		memdelete_arr(p_version->variants);
-		if (p_version->variant_data) {
-			memdelete_arr(p_version->variant_data);
-		}
-		p_version->variants = nullptr;
-		p_version->variant_data = nullptr;
+
+		p_version->variants.clear();
+		p_version->variant_data.clear();
 		return;
-	} else if (shader_cache_dir_valid) {
-		// Save shader cache.
+	}
+#if ENABLE_SHADER_CACHE
+	else if (shader_cache_dir_valid) {
 		_save_to_cache(p_version, p_group);
 	}
-
-	memdelete_arr(p_version->variant_data); //clear stages
-	p_version->variant_data = nullptr;
+#endif
 
 	p_version->valid = true;
 }
 
+void ShaderRD::_compile_ensure_finished(Version *p_version) {
+	// Wait for compilation of existing groups if necessary.
+	for (int i = 0; i < group_enabled.size(); i++) {
+		_compile_version_end(p_version, i);
+	}
+}
+
 void ShaderRD::version_set_code(RID p_version, const HashMap<String, String> &p_code, const String &p_uniforms, const String &p_vertex_globals, const String &p_fragment_globals, const Vector<String> &p_custom_defines) {
 	ERR_FAIL_COND(is_compute);
 
 	Version *version = version_owner.get_or_null(p_version);
 	ERR_FAIL_NULL(version);
+
+	_compile_ensure_finished(version);
+
 	version->vertex_globals = p_vertex_globals.utf8();
 	version->fragment_globals = p_fragment_globals.utf8();
 	version->uniforms = p_uniforms.utf8();
@@ -601,7 +613,7 @@ void ShaderRD::version_set_code(RID p_version, const HashMap<String, String> &p_
 				_allocate_placeholders(version, i);
 				continue;
 			}
-			_compile_version(version, i);
+			_compile_version_start(version, i);
 		}
 		version->initialize_needed = false;
 	}
@@ -613,6 +625,8 @@ void ShaderRD::version_set_compute_code(RID p_version, const HashMap<String, Str
 	Version *version = version_owner.get_or_null(p_version);
 	ERR_FAIL_NULL(version);
 
+	_compile_ensure_finished(version);
+
 	version->compute_globals = p_compute_globals.utf8();
 	version->uniforms = p_uniforms.utf8();
 
@@ -634,7 +648,7 @@ void ShaderRD::version_set_compute_code(RID p_version, const HashMap<String, Str
 				_allocate_placeholders(version, i);
 				continue;
 			}
-			_compile_version(version, i);
+			_compile_version_start(version, i);
 		}
 		version->initialize_needed = false;
 	}
@@ -651,10 +665,12 @@ bool ShaderRD::version_is_valid(RID p_version) {
 				_allocate_placeholders(version, i);
 				continue;
 			}
-			_compile_version(version, i);
+			_compile_version_start(version, i);
 		}
 	}
 
+	_compile_ensure_finished(version);
+
 	return version->valid;
 }
 
@@ -696,7 +712,7 @@ void ShaderRD::enable_group(int p_group) {
 	version_owner.get_owned_list(&all_versions);
 	for (const RID &E : all_versions) {
 		Version *version = version_owner.get_or_null(E);
-		_compile_version(version, p_group);
+		_compile_version_start(version, p_group);
 	}
 }
 
@@ -735,6 +751,7 @@ void ShaderRD::initialize(const Vector<String> &p_variant_defines, const String
 	for (int i = 0; i < p_variant_defines.size(); i++) {
 		variant_defines.push_back(VariantDefine(0, p_variant_defines[i], true));
 		variants_enabled.push_back(true);
+		variant_to_group.push_back(0);
 		group_to_variant_map[0].push_back(i);
 	}
 
@@ -796,6 +813,7 @@ void ShaderRD::initialize(const Vector<VariantDefine> &p_variant_defines, const
 		// Fill variant array.
 		variant_defines.push_back(p_variant_defines[i]);
 		variants_enabled.push_back(true);
+		variant_to_group.push_back(p_variant_defines[i].group);
 
 		// Map variant array index to group id, so we can iterate over groups later.
 		if (!group_to_variant_map.has(p_variant_defines[i].group)) {