Using a visitor to systematically replace lookup scopes of generic function's return type expression.

fixes #336

Add a `ReplaceScopeVisitor` to replace the scopes of the return type expression tree to use the generic decl's scope instead of module's scope after parser has determined the decl is a generic function header decl.

3 files changed, 70 insertions, 2 deletions

diff --git a/source/slang/parser.cpp b/source/slang/parser.cpp
index 35cc96b5c..662b75a2c 100644
--- a/source/slang/parser.cpp
+++ b/source/slang/parser.cpp
@@ -4,6 +4,7 @@
 
 #include "compiler.h"
 #include "lookup.h"
+#include "visitor.h"
 
 namespace Slang
 {
@@ -1099,6 +1100,40 @@ namespace Slang
         }
     }
 
+    // systematically replace all scopes in an expression tree
+    class ReplaceScopeVisitor : public ExprVisitor<ReplaceScopeVisitor>
+    {
+    public:
+        RefPtr<Scope> scope;
+        void visitDeclRefExpr(DeclRefExpr* expr) 
+        {
+            expr->scope = scope;
+        }
+        void visitGenericAppExpr(GenericAppExpr * expr)
+        {
+            expr->FunctionExpr->accept(this, nullptr);
+            for (auto arg : expr->Arguments)
+                arg->accept(this, nullptr);
+        }
+        void visitIndexExpr(IndexExpr * expr)
+        {
+            expr->BaseExpression->accept(this, nullptr);
+            expr->IndexExpression->accept(this, nullptr);
+        }
+        void visitMemberExpr(MemberExpr * expr)
+        {
+            expr->BaseExpression->accept(this, nullptr);
+            expr->scope = scope;
+        }
+        void visitStaticMemberExpr(StaticMemberExpr * expr)
+        {
+            expr->BaseExpression->accept(this, nullptr);
+            expr->scope = scope;
+        }
+        void visitExpr(Expr* /*expr*/)
+        {}
+    };
+
     static RefPtr<Decl> ParseFuncDeclHeader(
         Parser*                     parser,
         DeclaratorInfo const&       declaratorInfo,
@@ -1114,8 +1149,9 @@ namespace Slang
 
         // if return type is a DeclRef type, we need to update its scope to use this function decl's scope
         // so that LookUp can find the generic type parameters declared after the function name
-        if (auto declRefRetType = declaratorInfo.typeSpec.As<DeclRefExpr>())
-            declRefRetType->scope = parser->currentScope;
+        ReplaceScopeVisitor replaceScopeVisitor;
+        replaceScopeVisitor.scope = parser->currentScope;
+        declaratorInfo.typeSpec->accept(&replaceScopeVisitor, nullptr);
 
         decl->ReturnType = TypeExp(declaratorInfo.typeSpec);
         auto parseFuncDeclHeaderInner = [&](GenericDecl *)
diff --git a/tests/compute/generics-syntax-2.slang b/tests/compute/generics-syntax-2.slang
new file mode 100644
index 000000000..29ed64825
--- /dev/null
+++ b/tests/compute/generics-syntax-2.slang
@@ -0,0 +1,28 @@
+//TEST(compute):COMPARE_COMPUTE:-xslang -use-ir
+//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):dxbinding(0),glbinding(0),out
+// Confirm that generics syntax can be used in user
+// code and generates valid output.
+
+RWStructuredBuffer<float> outputBuffer;
+
+struct GenStruct<T>
+{
+    T x;
+};
+
+GenStruct<T> test<T>(T val)
+{
+    GenStruct<T> rs;
+    rs.x = val;
+    return rs;
+}
+
+
+[numthreads(4, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint tid = dispatchThreadID.x;
+    float inVal = float(tid);
+    GenStruct<float> outVal = test(inVal);
+    outputBuffer[tid] = outVal.x;
+}
\ No newline at end of file
diff --git a/tests/compute/generics-syntax-2.slang.expected.txt b/tests/compute/generics-syntax-2.slang.expected.txt
new file mode 100644
index 000000000..98798bd61
--- /dev/null
+++ b/tests/compute/generics-syntax-2.slang.expected.txt
@@ -0,0 +1,4 @@
+0
+3F800000
+40000000
+40400000
\ No newline at end of file