From e0374d735d406da9311e400a171445eaa13026ae Mon Sep 17 00:00:00 2001
From: Hans-Kristian Arntzen <post@arntzen-software.no>
Date: Wed, 17 Mar 2021 15:17:35 +0100
Subject: [PATCH] vkd3d-shader: Add shader replacement support for DXR as well.

Signed-off-by: Hans-Kristian Arntzen <post@arntzen-software.no>
---
 libs/vkd3d-shader/dxil.c                 | 17 +++++---
 libs/vkd3d-shader/vkd3d_shader_main.c    | 55 +++++++++++++++++-------
 libs/vkd3d-shader/vkd3d_shader_private.h |  1 +
 3 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/libs/vkd3d-shader/dxil.c b/libs/vkd3d-shader/dxil.c
index 7ee1afa8..19535da5 100644
--- a/libs/vkd3d-shader/dxil.c
+++ b/libs/vkd3d-shader/dxil.c
@@ -773,8 +773,16 @@ int vkd3d_shader_compile_dxil_export(const struct vkd3d_shader_code *dxil,
     hash = vkd3d_shader_hash(dxil);
     spirv->meta.replaced = false;
     spirv->meta.hash = hash;
-
-    /* TODO: Probably need to add per-export shader replacement. */
+    demangled_export = vkd3d_dup_demangled_entry_point_ascii(export);
+    if (demangled_export)
+    {
+        if (vkd3d_shader_replace_export(hash, &spirv->code, &spirv->size, demangled_export))
+        {
+            spirv->meta.replaced = true;
+            vkd3d_free(demangled_export);
+            return ret;
+        }
+    }
 
     dxil_spv_begin_thread_allocator_context();
 
@@ -1077,17 +1085,14 @@ int vkd3d_shader_compile_dxil_export(const struct vkd3d_shader_code *dxil,
     spirv->code = code;
     spirv->size = compiled.size;
 
-    demangled_export = vkd3d_dup_demangled_entry_point_ascii(export);
     if (demangled_export)
-    {
         vkd3d_shader_dump_spirv_shader_export(hash, spirv, demangled_export);
-        vkd3d_free(demangled_export);
-    }
 
 end:
     dxil_spv_converter_free(converter);
     dxil_spv_parsed_blob_free(blob);
     dxil_spv_end_thread_allocator_context();
+    vkd3d_free(demangled_export);
     return ret;
 }
 
diff --git a/libs/vkd3d-shader/vkd3d_shader_main.c b/libs/vkd3d-shader/vkd3d_shader_main.c
index b6c39e71..1a144cba 100644
--- a/libs/vkd3d-shader/vkd3d_shader_main.c
+++ b/libs/vkd3d-shader/vkd3d_shader_main.c
@@ -40,25 +40,12 @@ static void vkd3d_shader_dump_blob(const char *path, vkd3d_shader_hash_t hash, c
     }
 }
 
-bool vkd3d_shader_replace(vkd3d_shader_hash_t hash, const void **data, size_t *size)
+static bool vkd3d_shader_replace_path(const char *filename, vkd3d_shader_hash_t hash, const void **data, size_t *size)
 {
-    static bool enabled = true;
-    char filename[1024];
     void *buffer = NULL;
-    const char *path;
     FILE *f = NULL;
     size_t len;
 
-    if (!enabled)
-        return false;
-
-    if (!(path = getenv("VKD3D_SHADER_OVERRIDE")))
-    {
-        enabled = false;
-        return false;
-    }
-
-    snprintf(filename, ARRAY_SIZE(filename), "%s/%016"PRIx64".spv", path, hash);
     if ((f = fopen(filename, "rb")))
     {
         if (fseek(f, 0, SEEK_END) < 0)
@@ -78,7 +65,7 @@ bool vkd3d_shader_replace(vkd3d_shader_hash_t hash, const void **data, size_t *s
 
     *data = buffer;
     *size = len;
-    WARN("Overriding shader hash %016"PRIx64" with alternative SPIR-V module!\n", hash);
+    INFO("Overriding shader hash %016"PRIx64" with alternative SPIR-V module from %s!\n", hash, filename);
     fclose(f);
     return true;
 
@@ -89,6 +76,44 @@ err:
     return false;
 }
 
+bool vkd3d_shader_replace(vkd3d_shader_hash_t hash, const void **data, size_t *size)
+{
+    static bool enabled = true;
+    char filename[1024];
+    const char *path;
+
+    if (!enabled)
+        return false;
+
+    if (!(path = getenv("VKD3D_SHADER_OVERRIDE")))
+    {
+        enabled = false;
+        return false;
+    }
+
+    snprintf(filename, ARRAY_SIZE(filename), "%s/%016"PRIx64".spv", path, hash);
+    return vkd3d_shader_replace_path(filename, hash, data, size);
+}
+
+bool vkd3d_shader_replace_export(vkd3d_shader_hash_t hash, const void **data, size_t *size, const char *export)
+{
+    static bool enabled = true;
+    char filename[1024];
+    const char *path;
+
+    if (!enabled)
+        return false;
+
+    if (!(path = getenv("VKD3D_SHADER_OVERRIDE")))
+    {
+        enabled = false;
+        return false;
+    }
+
+    snprintf(filename, ARRAY_SIZE(filename), "%s/%016"PRIx64".lib.%s.spv", path, hash, export);
+    return vkd3d_shader_replace_path(filename, hash, data, size);
+}
+
 void vkd3d_shader_dump_shader(vkd3d_shader_hash_t hash, const struct vkd3d_shader_code *shader, const char *ext)
 {
     static bool enabled = true;
diff --git a/libs/vkd3d-shader/vkd3d_shader_private.h b/libs/vkd3d-shader/vkd3d_shader_private.h
index 30e1797e..c94451a0 100644
--- a/libs/vkd3d-shader/vkd3d_shader_private.h
+++ b/libs/vkd3d-shader/vkd3d_shader_private.h
@@ -801,6 +801,7 @@ void vkd3d_shader_dump_spirv_shader_export(vkd3d_shader_hash_t hash, const struc
         const char *export);
 void vkd3d_shader_dump_shader(vkd3d_shader_hash_t hash, const struct vkd3d_shader_code *shader, const char *ext);
 bool vkd3d_shader_replace(vkd3d_shader_hash_t hash, const void **data, size_t *size);
+bool vkd3d_shader_replace_export(vkd3d_shader_hash_t hash, const void **data, size_t *size, const char *export);
 
 static inline enum vkd3d_component_type vkd3d_component_type_from_data_type(
         enum vkd3d_data_type data_type)