From 1bbcf25af514a9ae24f7006747177f2d1b3b7c0d Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Mon, 11 Mar 2024 14:42:14 -0700
Subject: Link-time specialization fixes. (#3734)

* Fix method synthesis logic for static differentiable methods.

* Support link-time constants in thread group size reflection.
---
 source/slang/slang.cpp | 72 +++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 71 insertions(+), 1 deletion(-)

(limited to 'source/slang/slang.cpp')

diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp
index fe3f8dfa5..69c0f0e14 100644
--- a/source/slang/slang.cpp
+++ b/source/slang/slang.cpp
@@ -2137,7 +2137,7 @@ Type* ComponentType::getTypeFromString(
     // the modules that were directly or
     // indirectly referenced.
     //
-    Scope* scope = _createScopeForLegacyLookup(astBuilder);
+    Scope* scope = _getOrCreateScopeForLegacyLookup(astBuilder);
 
     auto linkage = getLinkage();
 
@@ -2154,6 +2154,76 @@ Type* ComponentType::getTypeFromString(
     return type;
 }
 
+static void collectExportedConstantInContainer(
+    Dictionary<String, IntVal*>& dict,
+    ASTBuilder* builder,
+    ContainerDecl* containerDecl)
+{
+    for (auto m : containerDecl->members)
+    {
+        auto varMember = as<VarDeclBase>(m);
+        if (!varMember)
+            continue;
+        if (!varMember->val)
+            continue;
+        bool isExported = false;
+        bool isConst = true;
+        bool isExtern = false;
+        for (auto modifier : m->modifiers)
+        {
+            if (as<HLSLExportModifier>(modifier))
+                isExported = true;
+            if (as<ExternAttribute>(modifier) || as<ExternModifier>(modifier))
+            {
+                isExtern = true;
+                isExported = true;
+            }
+            if (as<ConstModifier>(modifier))
+                isConst = true;
+            if (isExported && isConst)
+                break;
+        }
+        if (isExported && isConst)
+        {
+            auto mangledName = getMangledName(builder, m);
+            if (isExtern && dict.containsKey(mangledName))
+                continue;
+            dict[mangledName] = varMember->val;
+        }
+    }
+
+    for (auto member : containerDecl->members)
+    {
+        if (as<NamespaceDecl>(member) || as<FileDecl>(member))
+        {
+            collectExportedConstantInContainer(dict, builder, (ContainerDecl*)member);
+        }
+    }
+}
+
+Dictionary<String, IntVal*>& ComponentType::getMangledNameToIntValMap()
+{
+    if (m_mapMangledNameToIntVal)
+    {
+        return *m_mapMangledNameToIntVal;
+    }
+    m_mapMangledNameToIntVal = std::make_unique<Dictionary<String, IntVal*>>();
+    auto astBuilder = getLinkage()->getASTBuilder();
+    SLANG_AST_BUILDER_RAII(astBuilder);
+    Scope* scope = _getOrCreateScopeForLegacyLookup(astBuilder);
+    for (; scope; scope = scope->nextSibling)
+    {
+        if (scope->containerDecl)
+            collectExportedConstantInContainer(*m_mapMangledNameToIntVal, astBuilder, scope->containerDecl);
+    }
+    return *m_mapMangledNameToIntVal;
+}
+
+ConstantIntVal* ComponentType::tryFoldIntVal(IntVal* intVal)
+{
+    return as<ConstantIntVal>(intVal->linkTimeResolve(getMangledNameToIntValMap()));
+}
+
 CompileRequestBase::CompileRequestBase(
     Linkage*        linkage,
     DiagnosticSink* sink)
-- 
cgit v1.2.3