From a51c8c17e07063065fe1780d648c3b4808f31d59 Mon Sep 17 00:00:00 2001
From: squidbus <175574877+squidbus@users.noreply.github.com>
Date: Fri, 24 Jan 2025 12:47:04 -0800
Subject: [PATCH] shader_recompiler: Fix image write swizzles. (#2236)

---
 .../ir/passes/resource_tracking_pass.cpp      |  4 +--
 src/video_core/amdgpu/types.h                 | 31 ++++++++++++++++---
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp b/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
index d94c5223a..43d125bf0 100644
--- a/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
+++ b/src/shader_recompiler/ir/passes/resource_tracking_pass.cpp
@@ -569,7 +569,7 @@ void PatchTextureBufferArgs(IR::Block& block, IR::Inst& inst, Info& info) {
     inst.SetArg(1, CalculateBufferAddress(ir, inst, info, buffer, 1U));
 
     if (inst.GetOpcode() == IR::Opcode::StoreBufferFormatF32) {
-        const auto swizzled = ApplySwizzle(ir, inst.Arg(2), buffer.DstSelect());
+        const auto swizzled = ApplySwizzle(ir, inst.Arg(2), buffer.DstSelect().Inverse());
         const auto converted =
             ApplyWriteNumberConversionVec4(ir, swizzled, buffer.GetNumberConversion());
         inst.SetArg(2, converted);
@@ -829,7 +829,7 @@ void PatchImageArgs(IR::Block& block, IR::Inst& inst, Info& info) {
             auto texel = inst.Arg(4);
             if (is_storage) {
                 // Storage image requires shader swizzle.
-                texel = ApplySwizzle(ir, texel, image.DstSelect());
+                texel = ApplySwizzle(ir, texel, image.DstSelect().Inverse());
             }
             const auto converted =
                 ApplyWriteNumberConversionVec4(ir, texel, image.GetNumberConversion());
diff --git a/src/video_core/amdgpu/types.h b/src/video_core/amdgpu/types.h
index 63e184cc5..b442b2f1e 100644
--- a/src/video_core/amdgpu/types.h
+++ b/src/video_core/amdgpu/types.h
@@ -200,10 +200,10 @@ enum class NumberConversion : u32 {
 };
 
 struct CompMapping {
-    CompSwizzle r : 3;
-    CompSwizzle g : 3;
-    CompSwizzle b : 3;
-    CompSwizzle a : 3;
+    CompSwizzle r;
+    CompSwizzle g;
+    CompSwizzle b;
+    CompSwizzle a;
 
     auto operator<=>(const CompMapping& other) const = default;
 
@@ -217,6 +217,15 @@ struct CompMapping {
         };
     }
 
+    [[nodiscard]] CompMapping Inverse() const {
+        CompMapping result{};
+        InverseSingle(result.r, CompSwizzle::Red);
+        InverseSingle(result.g, CompSwizzle::Green);
+        InverseSingle(result.b, CompSwizzle::Blue);
+        InverseSingle(result.a, CompSwizzle::Alpha);
+        return result;
+    }
+
 private:
     template <typename T>
     T ApplySingle(const std::array<T, 4>& data, const CompSwizzle swizzle) const {
@@ -237,6 +246,20 @@ private:
             UNREACHABLE();
         }
     }
+
+    void InverseSingle(CompSwizzle& dst, const CompSwizzle target) const {
+        if (r == target) {
+            dst = CompSwizzle::Red;
+        } else if (g == target) {
+            dst = CompSwizzle::Green;
+        } else if (b == target) {
+            dst = CompSwizzle::Blue;
+        } else if (a == target) {
+            dst = CompSwizzle::Alpha;
+        } else {
+            dst = CompSwizzle::Zero;
+        }
+    }
 };
 
 inline DataFormat RemapDataFormat(const DataFormat format) {