Add shader baker to project exporter.

Metal Support contributed by Migeran (https://migeran.com) and Stuart Carnie.

Co-authored-by: Stuart Carnie <stuart.carnie@gmail.com>
Co-authored-by: Gergely Kis <gergely.kis@migeran.com>

servers/rendering/renderer_rd/shader_rd.cpp

@@ -152,10 +152,6 @@ void ShaderRD::setup(const char *p_vertex_code, const char *p_fragment_code, con
 	tohash.append(GODOT_VERSION_NUMBER);
 	tohash.append("[GodotVersionHash]");
 	tohash.append(GODOT_VERSION_HASH);
-	tohash.append("[SpirvCacheKey]");
-	tohash.append(RenderingDevice::get_singleton()->shader_get_spirv_cache_key());
-	tohash.append("[BinaryCacheKey]");
-	tohash.append(RenderingDevice::get_singleton()->shader_get_binary_cache_key());
 	tohash.append("[Vertex]");
 	tohash.append(p_vertex_code ? p_vertex_code : "");
 	tohash.append("[Fragment]");
@@ -166,7 +162,7 @@ void ShaderRD::setup(const char *p_vertex_code, const char *p_fragment_code, con
 	base_sha256 = tohash.as_string().sha256_text();
 }
 
-RID ShaderRD::version_create() {
+RID ShaderRD::version_create(bool p_embedded) {
 	//initialize() was never called
 	ERR_FAIL_COND_V(group_to_variant_map.is_empty(), RID());
 
@@ -174,12 +170,22 @@ RID ShaderRD::version_create() {
 	version.dirty = true;
 	version.valid = false;
 	version.initialize_needed = true;
+	version.embedded = p_embedded;
 	version.variants.clear();
 	version.variant_data.clear();
+
 	version.mutex = memnew(Mutex);
 	RID rid = version_owner.make_rid(version);
-	MutexLock lock(versions_mutex);
-	version_mutexes.insert(rid, version.mutex);
+	{
+		MutexLock lock(versions_mutex);
+		version_mutexes.insert(rid, version.mutex);
+	}
+
+	if (p_embedded) {
+		MutexLock lock(shader_versions_embedded_set_mutex);
+		shader_versions_embedded_set.insert({ this, rid });
+	}
+
 	return rid;
 }
 
@@ -263,86 +269,49 @@ void ShaderRD::_build_variant_code(StringBuilder &builder, uint32_t p_variant, c
 	}
 }
 
+Vector<String> ShaderRD::_build_variant_stage_sources(uint32_t p_variant, CompileData p_data) {
+	if (!variants_enabled[p_variant]) {
+		return Vector<String>(); // Variant is disabled, return.
+	}
+
+	Vector<String> stage_sources;
+	stage_sources.resize(RD::SHADER_STAGE_MAX);
+
+	if (is_compute) {
+		// Compute stage.
+		StringBuilder builder;
+		_build_variant_code(builder, p_variant, p_data.version, stage_templates[STAGE_TYPE_COMPUTE]);
+		stage_sources.write[RD::SHADER_STAGE_COMPUTE] = builder.as_string();
+	} else {
+		{
+			// Vertex stage.
+			StringBuilder builder;
+			_build_variant_code(builder, p_variant, p_data.version, stage_templates[STAGE_TYPE_VERTEX]);
+			stage_sources.write[RD::SHADER_STAGE_VERTEX] = builder.as_string();
+		}
+
+		{
+			// Fragment stage.
+			StringBuilder builder;
+			_build_variant_code(builder, p_variant, p_data.version, stage_templates[STAGE_TYPE_FRAGMENT]);
+			stage_sources.write[RD::SHADER_STAGE_FRAGMENT] = builder.as_string();
+		}
+	}
+
+	return stage_sources;
+}
+
 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.
 	}
 
-	Vector<RD::ShaderStageSPIRVData> stages;
-
-	String error;
-	String current_source;
-	RD::ShaderStage current_stage = RD::SHADER_STAGE_VERTEX;
-	bool build_ok = true;
-
-	if (!is_compute) {
-		//vertex stage
-
-		StringBuilder builder;
-		_build_variant_code(builder, variant, p_data.version, stage_templates[STAGE_TYPE_VERTEX]);
-
-		current_source = builder.as_string();
-		RD::ShaderStageSPIRVData stage;
-		stage.spirv = RD::get_singleton()->shader_compile_spirv_from_source(RD::SHADER_STAGE_VERTEX, current_source, RD::SHADER_LANGUAGE_GLSL, &error);
-		if (stage.spirv.is_empty()) {
-			build_ok = false;
-		} else {
-			stage.shader_stage = RD::SHADER_STAGE_VERTEX;
-			stages.push_back(stage);
-		}
-	}
-
-	if (!is_compute && build_ok) {
-		//fragment stage
-		current_stage = RD::SHADER_STAGE_FRAGMENT;
-
-		StringBuilder builder;
-		_build_variant_code(builder, variant, p_data.version, stage_templates[STAGE_TYPE_FRAGMENT]);
-
-		current_source = builder.as_string();
-		RD::ShaderStageSPIRVData stage;
-		stage.spirv = RD::get_singleton()->shader_compile_spirv_from_source(RD::SHADER_STAGE_FRAGMENT, current_source, RD::SHADER_LANGUAGE_GLSL, &error);
-		if (stage.spirv.is_empty()) {
-			build_ok = false;
-		} else {
-			stage.shader_stage = RD::SHADER_STAGE_FRAGMENT;
-			stages.push_back(stage);
-		}
-	}
-
-	if (is_compute) {
-		//compute stage
-		current_stage = RD::SHADER_STAGE_COMPUTE;
-
-		StringBuilder builder;
-		_build_variant_code(builder, variant, p_data.version, stage_templates[STAGE_TYPE_COMPUTE]);
-
-		current_source = builder.as_string();
-
-		RD::ShaderStageSPIRVData stage;
-		stage.spirv = RD::get_singleton()->shader_compile_spirv_from_source(RD::SHADER_STAGE_COMPUTE, current_source, RD::SHADER_LANGUAGE_GLSL, &error);
-		if (stage.spirv.is_empty()) {
-			build_ok = false;
-		} else {
-			stage.shader_stage = RD::SHADER_STAGE_COMPUTE;
-			stages.push_back(stage);
-		}
-	}
-
-	if (!build_ok) {
-		ERR_PRINT("Error compiling " + String(current_stage == RD::SHADER_STAGE_COMPUTE ? "Compute " : (current_stage == RD::SHADER_STAGE_VERTEX ? "Vertex" : "Fragment")) + " shader, variant #" + itos(variant) + " (" + variant_defines[variant].text.get_data() + ").");
-		ERR_PRINT(error);
-
-#ifdef DEBUG_ENABLED
-		ERR_PRINT("code:\n" + current_source.get_with_code_lines());
-#endif
-		return;
-	}
-
-	Vector<uint8_t> shader_data = RD::get_singleton()->shader_compile_binary_from_spirv(stages, name + ":" + itos(variant));
+	Vector<String> variant_stage_sources = _build_variant_stage_sources(variant, p_data);
+	Vector<RD::ShaderStageSPIRVData> variant_stages = compile_stages(variant_stage_sources);
+	ERR_FAIL_COND(variant_stages.is_empty());
 
+	Vector<uint8_t> shader_data = RD::get_singleton()->shader_compile_binary_from_spirv(variant_stages, name + ":" + itos(variant));
 	ERR_FAIL_COND(shader_data.is_empty());
 
 	{
@@ -351,6 +320,20 @@ void ShaderRD::_compile_variant(uint32_t p_variant, CompileData p_data) {
 	}
 }
 
+Vector<String> ShaderRD::version_build_variant_stage_sources(RID p_version, int p_variant) {
+	Version *version = version_owner.get_or_null(p_version);
+	ERR_FAIL_NULL_V(version, Vector<String>());
+
+	if (version->dirty) {
+		_initialize_version(version);
+	}
+
+	CompileData compile_data;
+	compile_data.version = version;
+	compile_data.group = variant_to_group[p_variant];
+	return _build_variant_stage_sources(p_variant, compile_data);
+}
+
 RS::ShaderNativeSourceCode ShaderRD::version_get_native_source_code(RID p_version) {
 	Version *version = version_owner.get_or_null(p_version);
 	RS::ShaderNativeSourceCode source_code;
@@ -404,6 +387,13 @@ RS::ShaderNativeSourceCode ShaderRD::version_get_native_source_code(RID p_versio
 	return source_code;
 }
 
+String ShaderRD::version_get_cache_file_relative_path(RID p_version, int p_group, const String &p_api_name) {
+	Version *version = version_owner.get_or_null(p_version);
+	ERR_FAIL_NULL_V(version, String());
+
+	return _get_cache_file_relative_path(version, p_group, p_api_name);
+}
+
 String ShaderRD::_version_get_sha1(Version *p_version) const {
 	StringBuilder hash_build;
 
@@ -437,17 +427,31 @@ String ShaderRD::_version_get_sha1(Version *p_version) const {
 static const char *shader_file_header = "GDSC";
 static const uint32_t cache_file_version = 4;
 
-String ShaderRD::_get_cache_file_path(Version *p_version, int p_group) {
-	const String &sha1 = _version_get_sha1(p_version);
-	const String &api_safe_name = String(RD::get_singleton()->get_device_api_name()).validate_filename().to_lower();
-	const String &path = shader_cache_dir.path_join(name).path_join(group_sha256[p_group]).path_join(sha1) + "." + api_safe_name + ".cache";
-	return path;
+String ShaderRD::_get_cache_file_relative_path(Version *p_version, int p_group, const String &p_api_name) {
+	String sha1 = _version_get_sha1(p_version);
+	return name.path_join(group_sha256[p_group]).path_join(sha1) + "." + p_api_name + ".cache";
+}
+
+String ShaderRD::_get_cache_file_path(Version *p_version, int p_group, const String &p_api_name, bool p_user_dir) {
+	const String &shader_cache_dir = p_user_dir ? shader_cache_user_dir : shader_cache_res_dir;
+	String relative_path = _get_cache_file_relative_path(p_version, p_group, p_api_name);
+	return shader_cache_dir.path_join(relative_path);
 }
 
 bool ShaderRD::_load_from_cache(Version *p_version, int p_group) {
-	const String &path = _get_cache_file_path(p_version, p_group);
-	Ref<FileAccess> f = FileAccess::open(path, FileAccess::READ);
+	String api_safe_name = String(RD::get_singleton()->get_device_api_name()).validate_filename().to_lower();
+	Ref<FileAccess> f;
+	if (shader_cache_user_dir_valid) {
+		f = FileAccess::open(_get_cache_file_path(p_version, p_group, api_safe_name, true), FileAccess::READ);
+	}
+
 	if (f.is_null()) {
+		f = FileAccess::open(_get_cache_file_path(p_version, p_group, api_safe_name, false), FileAccess::READ);
+	}
+
+	if (f.is_null()) {
+		const String &sha1 = _version_get_sha1(p_version);
+		print_verbose(vformat("Shader cache miss for %s", name.path_join(group_sha256[p_group]).path_join(sha1)));
 		return false;
 	}
 
@@ -506,19 +510,14 @@ bool ShaderRD::_load_from_cache(Version *p_version, int p_group) {
 }
 
 void ShaderRD::_save_to_cache(Version *p_version, int p_group) {
-	ERR_FAIL_COND(!shader_cache_dir_valid);
-	const String &path = _get_cache_file_path(p_version, p_group);
+	ERR_FAIL_COND(!shader_cache_user_dir_valid);
+	String api_safe_name = String(RD::get_singleton()->get_device_api_name()).validate_filename().to_lower();
+	const String &path = _get_cache_file_path(p_version, p_group, api_safe_name, true);
 	Ref<FileAccess> f = FileAccess::open(path, FileAccess::WRITE);
 	ERR_FAIL_COND(f.is_null());
-	f->store_buffer((const uint8_t *)shader_file_header, 4);
-	f->store_32(cache_file_version); // File version.
-	uint32_t variant_count = group_to_variant_map[p_group].size();
-	f->store_32(variant_count); // Variant count.
-	for (uint32_t i = 0; i < variant_count; i++) {
-		int variant_id = group_to_variant_map[p_group][i];
-		f->store_32(p_version->variant_data[variant_id].size()); // Stage count.
-		f->store_buffer(p_version->variant_data[variant_id].ptr(), p_version->variant_data[variant_id].size());
-	}
+
+	PackedByteArray shader_cache_bytes = ShaderRD::save_shader_cache_bytes(group_to_variant_map[p_group], p_version->variant_data);
+	f->store_buffer(shader_cache_bytes);
 }
 
 void ShaderRD::_allocate_placeholders(Version *p_version, int p_group) {
@@ -543,10 +542,8 @@ void ShaderRD::_compile_version_start(Version *p_version, int p_group) {
 	p_version->dirty = false;
 
 #if ENABLE_SHADER_CACHE
-	if (shader_cache_dir_valid) {
-		if (_load_from_cache(p_version, p_group)) {
-			return;
-		}
+	if (_load_from_cache(p_version, p_group)) {
+		return;
 	}
 #endif
 
@@ -595,7 +592,7 @@ void ShaderRD::_compile_version_end(Version *p_version, int p_group) {
 		return;
 	}
 #if ENABLE_SHADER_CACHE
-	else if (shader_cache_dir_valid) {
+	else if (shader_cache_user_dir_valid) {
 		_save_to_cache(p_version, p_group);
 	}
 #endif
@@ -714,6 +711,11 @@ bool ShaderRD::version_free(RID p_version) {
 		}
 
 		Version *version = version_owner.get_or_null(p_version);
+		if (version->embedded) {
+			MutexLock lock(shader_versions_embedded_set_mutex);
+			shader_versions_embedded_set.erase({ this, p_version });
+		}
+
 		version->mutex->lock();
 		_clear_version(version);
 		version_owner.free(p_version);
@@ -737,6 +739,14 @@ bool ShaderRD::is_variant_enabled(int p_variant) const {
 	return variants_enabled[p_variant];
 }
 
+int64_t ShaderRD::get_variant_count() const {
+	return variants_enabled.size();
+}
+
+int ShaderRD::get_variant_to_group(int p_variant) const {
+	return variant_to_group[p_variant];
+}
+
 void ShaderRD::enable_group(int p_group) {
 	ERR_FAIL_INDEX(p_group, group_enabled.size());
 
@@ -760,6 +770,18 @@ bool ShaderRD::is_group_enabled(int p_group) const {
 	return group_enabled[p_group];
 }
 
+int64_t ShaderRD::get_group_count() const {
+	return group_enabled.size();
+}
+
+const LocalVector<int> &ShaderRD::get_group_to_variants(int p_group) const {
+	return group_to_variant_map[p_group];
+}
+
+const String &ShaderRD::get_name() const {
+	return name;
+}
+
 bool ShaderRD::shader_cache_cleanup_on_start = false;
 
 ShaderRD::ShaderRD() {
@@ -778,12 +800,12 @@ ShaderRD::ShaderRD() {
 	base_compute_defines = base_compute_define_text.ascii();
 }
 
-void ShaderRD::initialize(const Vector<String> &p_variant_defines, const String &p_general_defines, const Vector<RD::PipelineImmutableSampler> &r_immutable_samplers) {
-	immutable_samplers = r_immutable_samplers;
+void ShaderRD::initialize(const Vector<String> &p_variant_defines, const String &p_general_defines, const Vector<RD::PipelineImmutableSampler> &p_immutable_samplers) {
 	ERR_FAIL_COND(variant_defines.size());
 	ERR_FAIL_COND(p_variant_defines.is_empty());
 
 	general_defines = p_general_defines.utf8();
+	immutable_samplers = p_immutable_samplers;
 
 	// When initialized this way, there is just one group and its always enabled.
 	group_to_variant_map.insert(0, LocalVector<int>{});
@@ -796,13 +818,18 @@ void ShaderRD::initialize(const Vector<String> &p_variant_defines, const String
 		group_to_variant_map[0].push_back(i);
 	}
 
-	if (!shader_cache_dir.is_empty()) {
+	if (!shader_cache_user_dir.is_empty() || !shader_cache_res_dir.is_empty()) {
 		group_sha256.resize(1);
 		_initialize_cache();
 	}
 }
 
 void ShaderRD::_initialize_cache() {
+	shader_cache_user_dir_valid = !shader_cache_user_dir.is_empty();
+	if (!shader_cache_user_dir_valid) {
+		return;
+	}
+
 	for (const KeyValue<int, LocalVector<int>> &E : group_to_variant_map) {
 		StringBuilder hash_build;
 
@@ -819,34 +846,44 @@ void ShaderRD::_initialize_cache() {
 
 		group_sha256[E.key] = hash_build.as_string().sha256_text();
 
-		Ref<DirAccess> d = DirAccess::open(shader_cache_dir);
-		ERR_FAIL_COND(d.is_null());
-		if (d->change_dir(name) != OK) {
-			Error err = d->make_dir(name);
-			ERR_FAIL_COND(err != OK);
-			d->change_dir(name);
-		}
+		if (!shader_cache_user_dir.is_empty()) {
+			// Validate if it's possible to write to all the directories required by in the user directory.
+			Ref<DirAccess> d = DirAccess::open(shader_cache_user_dir);
+			if (d.is_null()) {
+				shader_cache_user_dir_valid = false;
+				ERR_FAIL_MSG(vformat("Unable to open shader cache directory at %s.", shader_cache_user_dir));
+			}
 
-		// Erase other versions?
-		if (shader_cache_cleanup_on_start) {
+			if (d->change_dir(name) != OK) {
+				Error err = d->make_dir(name);
+				if (err != OK) {
+					shader_cache_user_dir_valid = false;
+					ERR_FAIL_MSG(vformat("Unable to create shader cache directory %s at %s.", name, shader_cache_user_dir));
+				}
+
+				d->change_dir(name);
+			}
+
+			if (d->change_dir(group_sha256[E.key]) != OK) {
+				Error err = d->make_dir(group_sha256[E.key]);
+				if (err != OK) {
+					shader_cache_user_dir_valid = false;
+					ERR_FAIL_MSG(vformat("Unable to create shader cache directory %s/%s at %s.", name, group_sha256[E.key], shader_cache_user_dir));
+				}
+			}
 		}
-		//
-		if (d->change_dir(group_sha256[E.key]) != OK) {
-			Error err = d->make_dir(group_sha256[E.key]);
-			ERR_FAIL_COND(err != OK);
-		}
-		shader_cache_dir_valid = true;
 
 		print_verbose("Shader '" + name + "' (group " + itos(E.key) + ") SHA256: " + group_sha256[E.key]);
 	}
 }
 
 // Same as above, but allows specifying shader compilation groups.
-void ShaderRD::initialize(const Vector<VariantDefine> &p_variant_defines, const String &p_general_defines) {
+void ShaderRD::initialize(const Vector<VariantDefine> &p_variant_defines, const String &p_general_defines, const Vector<RD::PipelineImmutableSampler> &p_immutable_samplers) {
 	ERR_FAIL_COND(variant_defines.size());
 	ERR_FAIL_COND(p_variant_defines.is_empty());
 
 	general_defines = p_general_defines.utf8();
+	immutable_samplers = p_immutable_samplers;
 
 	int max_group_id = 0;
 
@@ -877,14 +914,38 @@ void ShaderRD::initialize(const Vector<VariantDefine> &p_variant_defines, const
 		}
 	}
 
-	if (!shader_cache_dir.is_empty()) {
+	if (!shader_cache_user_dir.is_empty()) {
 		group_sha256.resize(max_group_id + 1);
 		_initialize_cache();
 	}
 }
 
-void ShaderRD::set_shader_cache_dir(const String &p_dir) {
-	shader_cache_dir = p_dir;
+void ShaderRD::shaders_embedded_set_lock() {
+	shader_versions_embedded_set_mutex.lock();
+}
+
+const ShaderRD::ShaderVersionPairSet &ShaderRD::shaders_embedded_set_get() {
+	return shader_versions_embedded_set;
+}
+
+void ShaderRD::shaders_embedded_set_unlock() {
+	shader_versions_embedded_set_mutex.unlock();
+}
+
+void ShaderRD::set_shader_cache_user_dir(const String &p_dir) {
+	shader_cache_user_dir = p_dir;
+}
+
+const String &ShaderRD::get_shader_cache_user_dir() {
+	return shader_cache_user_dir;
+}
+
+void ShaderRD::set_shader_cache_res_dir(const String &p_dir) {
+	shader_cache_res_dir = p_dir;
+}
+
+const String &ShaderRD::get_shader_cache_res_dir() {
+	return shader_cache_res_dir;
 }
 
 void ShaderRD::set_shader_cache_save_compressed(bool p_enable) {
@@ -899,7 +960,78 @@ void ShaderRD::set_shader_cache_save_debug(bool p_enable) {
 	shader_cache_save_debug = p_enable;
 }
 
-String ShaderRD::shader_cache_dir;
+Vector<RD::ShaderStageSPIRVData> ShaderRD::compile_stages(const Vector<String> &p_stage_sources) {
+	RD::ShaderStageSPIRVData stage;
+	Vector<RD::ShaderStageSPIRVData> stages;
+	String error;
+	RD::ShaderStage compilation_failed_stage = RD::SHADER_STAGE_MAX;
+	bool compilation_failed = false;
+	for (int64_t i = 0; i < p_stage_sources.size() && !compilation_failed; i++) {
+		if (p_stage_sources[i].is_empty()) {
+			continue;
+		}
+
+		stage.spirv = RD::get_singleton()->shader_compile_spirv_from_source(RD::ShaderStage(i), p_stage_sources[i], RD::SHADER_LANGUAGE_GLSL, &error);
+		stage.shader_stage = RD::ShaderStage(i);
+		if (!stage.spirv.is_empty()) {
+			stages.push_back(stage);
+
+		} else {
+			compilation_failed_stage = RD::ShaderStage(i);
+			compilation_failed = true;
+		}
+	}
+
+	if (compilation_failed) {
+		ERR_PRINT("Error compiling " + String(compilation_failed_stage == RD::SHADER_STAGE_COMPUTE ? "Compute " : (compilation_failed_stage == RD::SHADER_STAGE_VERTEX ? "Vertex" : "Fragment")) + " shader.");
+		ERR_PRINT(error);
+
+#ifdef DEBUG_ENABLED
+		ERR_PRINT("code:\n" + p_stage_sources[compilation_failed_stage].get_with_code_lines());
+#endif
+
+		return Vector<RD::ShaderStageSPIRVData>();
+	} else {
+		return stages;
+	}
+}
+
+PackedByteArray ShaderRD::save_shader_cache_bytes(const LocalVector<int> &p_variants, const Vector<Vector<uint8_t>> &p_variant_data) {
+	uint32_t variant_count = p_variants.size();
+	PackedByteArray bytes;
+	int64_t total_size = 0;
+	total_size += 4 + sizeof(uint32_t) * 2;
+	for (uint32_t i = 0; i < variant_count; i++) {
+		total_size += sizeof(uint32_t) + p_variant_data[p_variants[i]].size();
+	}
+
+	bytes.resize(total_size);
+
+	uint8_t *bytes_ptr = bytes.ptrw();
+	memcpy(bytes_ptr, shader_file_header, 4);
+	bytes_ptr += 4;
+
+	*(uint32_t *)(bytes_ptr) = cache_file_version;
+	bytes_ptr += sizeof(uint32_t);
+
+	*(uint32_t *)(bytes_ptr) = variant_count;
+	bytes_ptr += sizeof(uint32_t);
+
+	for (uint32_t i = 0; i < variant_count; i++) {
+		int variant_id = p_variants[i];
+		*(uint32_t *)(bytes_ptr) = uint32_t(p_variant_data[variant_id].size());
+		bytes_ptr += sizeof(uint32_t);
+
+		memcpy(bytes_ptr, p_variant_data[variant_id].ptr(), p_variant_data[variant_id].size());
+		bytes_ptr += p_variant_data[variant_id].size();
+	}
+
+	DEV_ASSERT((bytes.ptrw() + bytes.size()) == bytes_ptr);
+	return bytes;
+}
+
+String ShaderRD::shader_cache_user_dir;
+String ShaderRD::shader_cache_res_dir;
 bool ShaderRD::shader_cache_save_compressed = true;
 bool ShaderRD::shader_cache_save_compressed_zstd = true;
 bool ShaderRD::shader_cache_save_debug = true;