Support extension on generic type. (#4968)

4 files changed, 119 insertions, 1 deletions

diff --git a/tests/language-feature/extensions/generic-extension-1.slang b/tests/language-feature/extensions/generic-extension-1.slang
new file mode 100644
index 000000000..47940e31a
--- /dev/null
+++ b/tests/language-feature/extensions/generic-extension-1.slang
@@ -0,0 +1,49 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -shaderobj
+interface IFoo
+{
+    int getVal();
+}
+
+interface IBar
+{
+    int getValPlusOne();
+}
+
+interface IBaz
+{
+    int getValPlusTwo();
+}
+
+struct MyInt { int v; }
+
+extension MyInt : IFoo
+{
+    int getVal() { return v; }
+}
+
+// Since MyInt:IFoo, the following extension will make MyInt:IBar.
+extension<T: IFoo> T : IBar
+{
+    int getValPlusOne() { return this.getVal() + 1; }
+}
+
+// Since MyInt:IBar, the following extension will make MyInt:IBaz.
+extension<T: IBar> T : IBaz
+{
+    int getValPlusTwo() { return this.getValPlusOne() + 1; }
+}
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    MyInt v = {1};
+
+    // Check that the extensions applied to MyInt correctly, i.e.
+    // MyInt.getValPlusTwo() eixsts.
+
+    // CHECK: 3
+    outputBuffer[0] = v.getValPlusTwo();
+}
\ No newline at end of file
diff --git a/tests/language-feature/extensions/generic-extension-2.slang b/tests/language-feature/extensions/generic-extension-2.slang
new file mode 100644
index 000000000..2728a73d6
--- /dev/null
+++ b/tests/language-feature/extensions/generic-extension-2.slang
@@ -0,0 +1,30 @@
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -shaderobj
+interface IFoo<T>
+{
+    T getFirst();    
+}
+
+__generic<S, T : IArray<S>>
+extension T : IFoo<S>
+{ 
+    S getFirst()
+    {
+        return this[0];
+    }
+}
+
+T getFirstOuter<T>(IFoo<T> arr)
+{
+    return arr.getFirst();
+}
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(1,1,1)]
+void computeMain()
+{
+    int arr[] = {1,2,3};
+    // CHECK: 1
+    outputBuffer[0] = getFirstOuter(arr);
+}
\ No newline at end of file
diff --git a/tests/language-feature/extensions/interface-extension.slang b/tests/language-feature/extensions/interface-extension.slang
index 1d84ba844..50bbbe22a 100644
--- a/tests/language-feature/extensions/interface-extension.slang
+++ b/tests/language-feature/extensions/interface-extension.slang
@@ -16,7 +16,8 @@ struct MyCounter : ICounter
 	[mutating] void add(int value) { _state += value; }
 }
 
-extension ICounter
+__generic<T : ICounter>
+extension T
 {
 	[mutating] void increment()
 	{
diff --git a/tests/language-feature/extensions/this-in-extension.slang b/tests/language-feature/extensions/this-in-extension.slang
new file mode 100644
index 000000000..374eabe6f
--- /dev/null
+++ b/tests/language-feature/extensions/this-in-extension.slang
@@ -0,0 +1,38 @@
+// this-in-extension.slang
+
+// Test that an `This` type works correctly when there is an extension.
+
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK): -shaderobj
+
+interface IFoo
+{
+    static const This identity;
+}
+
+__generic<T:IFoo>
+extension T
+{
+    This getIdentity()
+    {
+        return identity;
+    }
+}
+
+struct FooImpl : IFoo
+{
+    int v = 1;
+    static const This identity = This();   
+}
+
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[numthreads(1, 1, 1)]
+void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
+{
+    FooImpl foo;
+    var ident = foo.getIdentity();
+    // CHECK: 1
+    outputBuffer[0] = ident.v;
+}