From cdd5e6666f98903d61118ada5cba51424fd1577c Mon Sep 17 00:00:00 2001
From: Ellie Hermaszewska <ellieh@nvidia.com>
Date: Wed, 23 Aug 2023 17:29:25 +0800
Subject: Retain int casts when unifying generic params (#3145)

---
 source/slang/slang-check-constraint.cpp          | 40 ++++++++++--------------
 tests/bugs/generic-param-cast.slang              | 37 ++++++++++++++++++++++
 tests/bugs/generic-param-cast.slang.expected.txt |  4 +++
 3 files changed, 58 insertions(+), 23 deletions(-)
 create mode 100644 tests/bugs/generic-param-cast.slang
 create mode 100644 tests/bugs/generic-param-cast.slang.expected.txt

diff --git a/source/slang/slang-check-constraint.cpp b/source/slang/slang-check-constraint.cpp
index e7eccfaec..6e600c4af 100644
--- a/source/slang/slang-check-constraint.cpp
+++ b/source/slang/slang-check-constraint.cpp
@@ -526,30 +526,24 @@ namespace Slang
         }
 
         // Check if both are integer values in general
-        if (auto fstInt = as<IntVal>(fst))
+        const auto fstInt = as<IntVal>(fst);
+        const auto sndInt = as<IntVal>(snd);
+        if (fstInt && sndInt)
         {
-            if (auto tc = as<TypeCastIntVal>(fstInt))
-                fstInt = as<IntVal>(tc->getBase());
-            if (auto sndInt = as<IntVal>(snd))
-            {
-                if (auto tc = as<TypeCastIntVal>(sndInt))
-                    sndInt = as<IntVal>(tc->getBase());
-                auto fstParam = as<GenericParamIntVal>(fstInt);
-                auto sndParam = as<GenericParamIntVal>(sndInt);
-
-                bool okay = false;
-                if (fstParam)
-                {
-                    if(TryUnifyIntParam(constraints, fstParam->getDeclRef(), sndInt))
-                        okay = true;
-                }
-                if (sndParam)
-                {
-                    if(TryUnifyIntParam(constraints, sndParam->getDeclRef(), fstInt))
-                        okay = true;
-                }
-                return okay;
-            }
+            const auto paramUnderCast = [](IntVal* i){
+                if(const auto c = as<TypeCastIntVal>(i))
+                    i = as<IntVal>(c->getBase());
+                return as<GenericParamIntVal>(i);
+            };
+            auto fstParam = paramUnderCast(fstInt);
+            auto sndParam = paramUnderCast(sndInt);
+
+            bool okay = false;
+            if (fstParam)
+                okay |= TryUnifyIntParam(constraints, fstParam->getDeclRef(), sndInt);
+            if (sndParam)
+                okay |= TryUnifyIntParam(constraints, sndParam->getDeclRef(), fstInt);
+            return okay;
         }
 
         if (auto fstWit = as<DeclaredSubtypeWitness>(fst))
diff --git a/tests/bugs/generic-param-cast.slang b/tests/bugs/generic-param-cast.slang
new file mode 100644
index 000000000..20b30a433
--- /dev/null
+++ b/tests/bugs/generic-param-cast.slang
@@ -0,0 +1,37 @@
+//TEST(smoke,compute):COMPARE_COMPUTE:-cpu -shaderobj
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+struct A<let I : int>
+{
+    int f() { return I; }
+};
+
+struct B<let U : uint>
+{
+    A<U> a;
+};
+
+int foo<let I : int>(A<I> a)
+{
+    return a.f();
+}
+
+int bar<let U : uint>(B<U> b)
+{
+    return foo(b.a);
+    // We previously were inferring the type at which to call `foo` as `U`
+    // instead of `int(U)`. This then cause the typechecker to impmediately
+    // fail because `U` does not unify with the type of `B<U>`'s `a` member,
+    // namely `int(U)`.
+}
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    A<1> a;
+    B<1> b;
+    b.a = a;
+    outputBuffer[dispatchThreadID.x] = bar<1>(b);
+}
diff --git a/tests/bugs/generic-param-cast.slang.expected.txt b/tests/bugs/generic-param-cast.slang.expected.txt
new file mode 100644
index 000000000..98fb6a686
--- /dev/null
+++ b/tests/bugs/generic-param-cast.slang.expected.txt
@@ -0,0 +1,4 @@
+1
+1
+1
+1
-- 
cgit v1.2.3