From 2ba53e06faa228a1aaee46240be2fe4d743ce22b Mon Sep 17 00:00:00 2001
From: Shaun Ren <sren@codeweavers.com>
Date: Thu, 16 Oct 2025 21:57:29 -0400
Subject: [PATCH] tests/hlsl: Add some conditional flattening tests.

---
 tests/hlsl/conditional.shader_test | 130 +++++++++++++++++++++++++++++
 tests/hlsl/discard.shader_test     |  62 ++++++++++++++
 2 files changed, 192 insertions(+)

diff --git a/tests/hlsl/conditional.shader_test b/tests/hlsl/conditional.shader_test
index d4c6b432d..607fcb065 100644
--- a/tests/hlsl/conditional.shader_test
+++ b/tests/hlsl/conditional.shader_test
@@ -110,6 +110,35 @@ uniform 0 float4 0.0 0.0 0.0 0.0
 draw quad
 probe (0, 0) rgba (0.9, 0.8, 0.7, 0.6)
 
+[pixel shader]
+uniform float4 u;
+
+float4 main() : sv_target
+{
+    [flatten]
+    if (u.x > 0.0)
+    {
+        [branch]
+        if (u.y > 0.0)
+            return float4(1.0, 2.0, 3.0, 4.0);
+        else
+            return float4(5.0, 6.0, 7.0, 8.0);
+    }
+    else
+        return float4(9.0, 10.0, 11.0, 12.0);
+}
+
+[test]
+uniform 0 float4 1.0 0.0 0.0 0.0
+draw quad
+probe (0, 0) f32(5.0, 6.0, 7.0, 8.0)
+uniform 0 float4 1.0 1.0 0.0 0.0
+draw quad
+probe (0, 0) f32(1.0, 2.0, 3.0, 4.0)
+uniform 0 float4 0.0 1.0 0.0 0.0
+draw quad
+probe (0, 0) f32(9.0, 10.0, 11.0, 12.0)
+
 [pixel shader]
 float4 main() : sv_target
 {
@@ -215,6 +244,54 @@ if(sm>=4) uniform 1 int 4
 draw quad
 probe (0, 0) rgba (1.0, 1.0, 1.0, 1.0)
 
+[pixel shader]
+uniform float4 u;
+
+float4 main() : sv_target
+{
+    float res = 0;
+
+    [flatten]
+    if (u.x > 0.0)
+        res = 1.0;
+    else
+        res = 2.0;
+
+    return res;
+}
+
+[test]
+uniform 0 float4 0.0 0.0 0.0 0.0
+draw quad
+probe (0, 0) f32(2.0, 2.0, 2.0, 2.0)
+uniform 0 float4 1.0 0.0 0.0 0.0
+draw quad
+probe (0, 0) f32(1.0, 1.0, 1.0, 1.0)
+
+[pixel shader]
+uniform float4 u;
+
+float4 main() : sv_target
+{
+    float4 res = float4(0.0, -1.0, -2.0, -3.0);
+
+    [flatten]
+    if (u.x > 0.0)
+        res.xy = float2(4.0, 5.0);
+    else
+        res.xz = float2(6.0, 7.0);
+
+    return res;
+}
+
+[test]
+uniform 0 float4 0.0 0.0 0.0 0.0
+draw quad
+probe (0, 0) f32(6.0, -1.0, 7.0, -3.0)
+uniform 0 float4 1.0 0.0 0.0 0.0
+draw quad
+probe (0, 0) f32(4.0, 5.0, -2.0, -3.0)
+
 % Test "if" conditionals, using resource loads to ensure we aren't removing these instruction with optimizations.
 [srv 0]
 size (2d, 2, 2)
@@ -250,6 +327,27 @@ Texture2D tex;
 sampler sam;
 float a;
 
+float4 main() : sv_target
+{
+    [flatten] if (a < 4)
+        return tex.Sample(sam, float2(0, 0));
+    else
+        return float4(1, 2, 3, 4);
+}
+
+[test]
+uniform 0 float -2
+draw quad
+probe (0, 0) f32(0.0, 0.0, 0.0, 4.0)
+uniform 0 float 4
+draw quad
+probe (0, 0) f32(1.0, 2.0, 3.0, 4.0)
+
+[pixel shader]
+Texture2D tex;
+sampler sam;
+float a;
+
 float4 main() : sv_target
 {
     if (a >= 2)
@@ -271,6 +369,38 @@ uniform 0 float 0
 todo(msl & sm>=6) draw quad
 probe (0, 0) rgba (1.0, 1.0, 0.0, 4.0)
 
+% Branches with resource stores can't be flattened.
+[pixel shader fail(sm<6) todo]
+uniform float a;
+RWBuffer<float> u;
+
+float4 main() : sv_target
+{
+    [flatten]
+    if (a > 0)
+       u[0] = a;
+    return a;
+}
+
+[pixel shader fail(sm<6) todo]
+uniform float a;
+uniform uint b;
+RWBuffer<float> u;
+
+float4 main() : sv_target
+{
+    [flatten]
+    if (a > 0)
+    {
+       for (uint i = 0; i < b; ++i)
+       {
+           if (i > 3)
+               u[i] = a + i;
+       }
+    }
+    return a;
+}
+
 % Test 16-bit phi instructions.
 [require]
 shader model >= 6.2
diff --git a/tests/hlsl/discard.shader_test b/tests/hlsl/discard.shader_test
index 999415bda..8ed9ceeae 100644
--- a/tests/hlsl/discard.shader_test
+++ b/tests/hlsl/discard.shader_test
@@ -16,6 +16,68 @@ uniform 0 float4 9 8 7 6
 todo(sm<4 | glsl | msl & sm>=6) draw quad
 probe (0, 0) f32(1, 2, 3, 4)
 
+[pixel shader todo(sm<4)]
+uniform float4 x;
+
+float4 main() : sv_target
+{
+    [flatten]
+    if (x.x == 9.0f)
+        discard;
+    return x;
+}
+
+[test]
+uniform 0 float4 1 2 3 4
+todo(sm<4 | glsl | msl & sm>=6) draw quad
+probe (0, 0) f32(1, 2, 3, 4)
+uniform 0 float4 9 8 7 6
+todo(sm<4 | glsl | msl & sm>=6) draw quad
+probe (0, 0) f32(1, 2, 3, 4)
+
+[pixel shader todo(sm<4)]
+uniform float4 x;
+
+float4 main() : sv_target
+{
+    float4 ret = x;
+    if (x.x != 9.0f)
+        ret.w = 4.0f;
+    else
+        discard;
+    return ret;
+}
+
+[test]
+uniform 0 float4 1 2 3 4
+todo(sm<4 | glsl | msl & sm>=6) draw quad
+probe (0, 0) f32(1, 2, 3, 4)
+uniform 0 float4 9 8 7 6
+todo(sm<4 | glsl | msl & sm>=6) draw quad
+probe (0, 0) f32(1, 2, 3, 4)
+
+[pixel shader todo(sm<4)]
+uniform float4 x;
+
+float4 main() : sv_target
+{
+    float4 ret = x;
+    [flatten]
+    if (x.x != 9.0f)
+        ret.w = 4.0f;
+    else
+        discard;
+    return ret;
+}
+
+[test]
+uniform 0 float4 1 2 3 4
+todo(sm<4 | glsl | msl & sm>=6) draw quad
+probe (0, 0) f32(1, 2, 3, 4)
+uniform 0 float4 9 8 7 6
+todo(sm<4 | glsl | msl & sm>=6) draw quad
+probe (0, 0) f32(1, 2, 3, 4)
+
 [require]
 shader model >= 3.0