From a0bea07503c68160ad2e88986ba98cfc2161bdff Mon Sep 17 00:00:00 2001
From: Sai Praveen Bangaru <31557731+saipraveenb25@users.noreply.github.com>
Date: Wed, 23 Oct 2024 20:28:49 -0400
Subject: Fix several bugs with `specializeWithArgTypes()` (#5365)

* Fix several bugs with `specializeWithArgTypes()`

* Make all types L-values for the purposes of reflection API resolution
---
 source/slang/slang-compiler.h         |  2 +
 source/slang/slang-reflection-api.cpp | 13 ++++--
 source/slang/slang.cpp                | 80 ++++++++++++++++++++++++++++++-----
 3 files changed, 81 insertions(+), 14 deletions(-)

(limited to 'source')

diff --git a/source/slang/slang-compiler.h b/source/slang/slang-compiler.h
index 980490df5..c1251488b 100755
--- a/source/slang/slang-compiler.h
+++ b/source/slang/slang-compiler.h
@@ -2281,6 +2281,8 @@ namespace Slang
             Expr*               funcExpr,
             List<Type*>         argTypes,
             DiagnosticSink*     sink);
+        
+        bool isSpecialized(DeclRef<Decl> declRef);
 
         DiagnosticSink::Flags diagnosticSinkFlags = 0;
 
diff --git a/source/slang/slang-reflection-api.cpp b/source/slang/slang-reflection-api.cpp
index 45bbd400c..ae351aee9 100644
--- a/source/slang/slang-reflection-api.cpp
+++ b/source/slang/slang-reflection-api.cpp
@@ -808,9 +808,9 @@ SlangReflectionFunction* tryConvertExprToFunctionReflection(ASTBuilder* astBuild
         auto declRef = declRefExpr->declRef;
         if (auto genericDeclRef = declRef.as<GenericDecl>())
         {
-            auto innerDeclRef = substituteDeclRef(
-                SubstitutionSet(genericDeclRef), astBuilder, genericDeclRef.getDecl()->inner);
-            declRef = createDefaultSubstitutionsIfNeeded(astBuilder, nullptr, innerDeclRef);
+            auto innerDeclRef = createDefaultSubstitutionsIfNeeded(astBuilder, nullptr, genericDeclRef.getDecl()->inner);
+            declRef = substituteDeclRef(
+                SubstitutionSet(genericDeclRef), astBuilder, innerDeclRef);
         }
 
         if (auto funcDeclRef = declRef.as<FunctionDeclBase>())
@@ -3194,7 +3194,12 @@ SLANG_API SlangReflectionFunction* spReflectionFunction_specializeWithArgTypes(
     try 
     {
         DiagnosticSink sink(linkage->getSourceManager(), Lexer::sourceLocationLexer);
-        return convert(linkage->specializeWithArgTypes(funcExpr, argTypeList, &sink).as<FunctionDeclBase>());
+        auto resultFunc = linkage->specializeWithArgTypes(funcExpr, argTypeList, &sink).as<FunctionDeclBase>();
+
+        if (sink.getErrorCount() != 0)
+            return nullptr; // Failed coercion.
+
+        return convert(resultFunc);
     }
     catch (...)
     {
diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index 313e94439..7949fd9f7 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -1369,16 +1369,68 @@ DeclRef<GenericDecl> getGenericParentDeclRef(
     // Create substituted parent decl ref.
     auto decl = declRef.getDecl();
 
-    while (!as<GenericDecl>(decl))
+    while (decl && !as<GenericDecl>(decl))
     {
         decl = decl->parentDecl;
     }
 
+    if (!decl)
+    {
+        // No generic parent
+        return DeclRef<GenericDecl>();
+    }
+
     auto genericDecl = as<GenericDecl>(decl);
     auto genericDeclRef = createDefaultSubstitutionsIfNeeded(astBuilder, visitor, DeclRef(genericDecl)).as<GenericDecl>();
     return substituteDeclRef(SubstitutionSet(declRef), astBuilder, genericDeclRef).as<GenericDecl>();
 }
 
+bool Linkage::isSpecialized(DeclRef<Decl> declRef)
+{
+    // For now, we only support two 'states': fully applied or not at all.
+    // If we add support for partial specialization, we will need to update this logic.
+    // 
+    // If it's not specialized, then declRef will be the one with default substitutions.
+    //
+    SemanticsVisitor visitor(getSemanticsForReflection());
+
+    auto decl = declRef.getDecl();
+    while (decl && !as<GenericDecl>(decl))
+    {
+        decl = decl->parentDecl;
+    }
+
+    if(!decl)
+        return true; // no generics => always specialized
+    
+    auto defaultArgs = getDefaultSubstitutionArgs(getASTBuilder(), &visitor, as<GenericDecl>(decl));
+    auto currentArgs = SubstitutionSet(declRef).findGenericAppDeclRef(as<GenericDecl>(decl))->getArgs();
+    
+    if (defaultArgs.getCount() != currentArgs.getCount()) // should really never happen.
+        return true;
+
+    for (Index i = 0; i < defaultArgs.getCount(); ++i)
+    {
+        if (defaultArgs[i] != currentArgs[i])
+            return true;
+    }
+
+    return false;
+}
+
+bool isFuncGeneric(DeclRef<Decl> declRef)
+{
+    if (auto funcDecl = as<FuncDecl>(declRef.getDecl()))
+    {
+        if (funcDecl->parentDecl && as<GenericDecl>(funcDecl->parentDecl))
+        {
+            return true;
+        }
+    }
+    
+    return false;
+}
+
 DeclRef<Decl> Linkage::specializeWithArgTypes(
     Expr*               funcExpr,
     List<Type*>         argTypes,
@@ -1387,16 +1439,22 @@ DeclRef<Decl> Linkage::specializeWithArgTypes(
     SemanticsVisitor visitor(getSemanticsForReflection());
     visitor = visitor.withSink(sink);
 
-    ASTBuilder* astBuilder = getASTBuilder();
+    SLANG_AST_BUILDER_RAII(getASTBuilder());
     
     if (auto declRefFuncExpr = as<DeclRefExpr>(funcExpr))
     {
-        auto genericDeclRefExpr = astBuilder->create<DeclRefExpr>();
-        genericDeclRefExpr->declRef = getGenericParentDeclRef(
-            getASTBuilder(),
-            &visitor,
-            declRefFuncExpr->declRef);
-        funcExpr = genericDeclRefExpr;
+        if (isFuncGeneric(declRefFuncExpr->declRef) && !isSpecialized(declRefFuncExpr->declRef))
+        {
+            if (auto genericDeclRef = getGenericParentDeclRef(
+                getCurrentASTBuilder(),
+                &visitor,
+                declRefFuncExpr->declRef))
+            {
+                auto genericDeclRefExpr = getCurrentASTBuilder()->create<DeclRefExpr>();
+                genericDeclRefExpr->declRef = genericDeclRef;
+                funcExpr = genericDeclRefExpr;
+            }
+        }
     }
 
     List<Expr*> argExprs;
@@ -1407,17 +1465,19 @@ DeclRef<Decl> Linkage::specializeWithArgTypes(
         // Create an 'empty' expr with the given type. Ideally, the expression itself should not matter
         // only its checked type.
         //
-        auto argExpr = astBuilder->create<VarExpr>();
+        auto argExpr = getCurrentASTBuilder()->create<VarExpr>();
         argExpr->type = argType;
+        argExpr->type.isLeftValue = true;
         argExprs.add(argExpr);
     }
 
     // Construct invoke expr.
-    auto invokeExpr = astBuilder->create<InvokeExpr>();
+    auto invokeExpr = getCurrentASTBuilder()->create<InvokeExpr>();
     invokeExpr->functionExpr = funcExpr;
     invokeExpr->arguments = argExprs;
 
     auto checkedInvokeExpr = visitor.CheckInvokeExprWithCheckedOperands(invokeExpr);
+
     return as<DeclRefExpr>(as<InvokeExpr>(checkedInvokeExpr)->functionExpr)->declRef;
 }
 
-- 
cgit v1.2.3