From 618b4c7657f539e66f032cd40554798bc0d68f6d Mon Sep 17 00:00:00 2001
From: Ellie Hermaszewska <ellieh@nvidia.com>
Date: Sat, 1 Mar 2025 05:22:58 +0800
Subject: Resolve 'extern' types during type layout generation if possible
 (#6450)

* Resolve 'extern' types during type layout generation if possible

Closes https://github.com/shader-slang/slang/issues/5994
Closes https://github.com/shader-slang/slang/issues/6437

* format code

---------

Co-authored-by: slangbot <186143334+slangbot@users.noreply.github.com>
Co-authored-by: Yong He <yonghe@outlook.com>
---
 source/slang/slang-parameter-binding.cpp |  7 +++
 source/slang/slang-type-layout.cpp       | 79 ++++++++++++++++++++++++++++++++
 source/slang/slang-type-layout.h         |  7 +++
 3 files changed, 93 insertions(+)

(limited to 'source')

diff --git a/source/slang/slang-parameter-binding.cpp b/source/slang/slang-parameter-binding.cpp
index 0e47ff56d..86d76fc6d 100644
--- a/source/slang/slang-parameter-binding.cpp
+++ b/source/slang/slang-parameter-binding.cpp
@@ -2382,7 +2382,14 @@ static RefPtr<TypeLayout> processEntryPointVaryingParameter(
     // otherwise they will include all of the above cases...
     else if (auto declRefType = as<DeclRefType>(type))
     {
+        // If we are trying to get the layout of some extern type, do our best
+        // to look it up in other loaded modules and generate the type layout
+        // based on that.
+        declRefType = context->layoutContext.lookupExternDeclRefType(declRefType);
+
         auto declRef = declRefType->getDeclRef();
+
+
         if (auto structDeclRef = declRef.as<StructDecl>())
         {
             RefPtr<StructTypeLayout> structLayout = new StructTypeLayout();
diff --git a/source/slang/slang-type-layout.cpp b/source/slang/slang-type-layout.cpp
index a412bf5b2..7968450b4 100644
--- a/source/slang/slang-type-layout.cpp
+++ b/source/slang/slang-type-layout.cpp
@@ -4,6 +4,7 @@
 #include "../compiler-core/slang-artifact-desc-util.h"
 #include "slang-check-impl.h"
 #include "slang-ir-insts.h"
+#include "slang-mangle.h"
 #include "slang-syntax.h"
 
 #include <assert.h>
@@ -5013,8 +5014,13 @@ static TypeLayoutResult _createTypeLayout(TypeLayoutContext& context, Type* type
     }
     else if (auto declRefType = as<DeclRefType>(type))
     {
+        // If we are trying to get the layout of some extern type, do our best
+        // to look it up in other loaded modules and generate the type layout
+        // based on that.
+        declRefType = context.lookupExternDeclRefType(declRefType);
         auto declRef = declRefType->getDeclRef();
 
+
         if (auto structDeclRef = declRef.as<StructDecl>())
         {
             StructTypeLayoutBuilder typeLayoutBuilder;
@@ -5694,4 +5700,77 @@ GlobalGenericParamDecl* GenericParamTypeLayout::getGlobalGenericParamDecl()
     return rsDeclRef.getDecl();
 }
 
+DeclRefType* TypeLayoutContext::lookupExternDeclRefType(DeclRefType* declRefType)
+{
+    const auto declRef = declRefType->getDeclRef();
+    const auto decl = declRef.getDecl();
+    const auto isExtern =
+        decl->hasModifier<ExternAttribute>() || decl->hasModifier<ExternModifier>();
+    if (isExtern)
+    {
+        if (!externTypeMap)
+            buildExternTypeMap();
+        const auto mangledName = getMangledName(targetReq->getLinkage()->getASTBuilder(), decl);
+        externTypeMap->tryGetValue(mangledName, declRefType);
+    }
+    return declRefType;
+}
+
+void TypeLayoutContext::buildExternTypeMap()
+{
+    externTypeMap.emplace();
+    const auto linkage = targetReq->getLinkage();
+
+    HashSet<String> externNames;
+    Dictionary<String, DeclRefType*> allTypes;
+
+    // Traverse the AST and keep track of all extern names and all type definitions
+    // We'll match them up later
+    auto processDecl = [&](auto&& go, Decl* decl) -> void
+    {
+        const auto isExtern =
+            decl->hasModifier<ExternAttribute>() || decl->hasModifier<ExternModifier>();
+
+        if (auto declRefType = as<DeclRefType>(DeclRefType::create(astBuilder, decl)))
+        {
+            String mangledName = getMangledName(astBuilder, decl);
+
+            if (isExtern)
+            {
+                externNames.add(mangledName);
+            }
+            else
+            {
+                allTypes[mangledName] = declRefType;
+            }
+        }
+
+        if (auto scopeDecl = as<ScopeDecl>(decl))
+        {
+            for (auto member : scopeDecl->members)
+            {
+                go(go, member);
+            }
+        }
+    };
+
+    for (const auto& m : linkage->loadedModulesList)
+    {
+        const auto& ast = m->getModuleDecl();
+        for (auto member : ast->members)
+        {
+            processDecl(processDecl, member);
+        }
+    }
+
+    // Only keep the types that have matching extern declarations
+    for (const auto& externName : externNames)
+    {
+        if (allTypes.containsKey(externName))
+        {
+            externTypeMap.value()[externName] = allTypes[externName];
+        }
+    }
+}
+
 } // namespace Slang
diff --git a/source/slang/slang-type-layout.h b/source/slang/slang-type-layout.h
index 7d3dd8369..bcb842013 100644
--- a/source/slang/slang-type-layout.h
+++ b/source/slang/slang-type-layout.h
@@ -1183,6 +1183,13 @@ struct TypeLayoutContext
     // Options passed to object layout
     ObjectLayoutRulesImpl::Options objectLayoutOptions;
 
+    // Mangled names to DeclRefType, this is used to match up 'extern' types to
+    // their linked in definitions during layout generation
+    std::optional<Dictionary<String, DeclRefType*>> externTypeMap;
+
+    DeclRefType* lookupExternDeclRefType(DeclRefType* declRefType);
+    void buildExternTypeMap();
+
     LayoutRulesImpl* getRules() { return rules; }
     LayoutRulesFamilyImpl* getRulesFamily() const { return rules->getLayoutRulesFamily(); }
 
-- 
cgit v1.2.3