Resolve 'extern' types during type layout generation if possible

Closes shader-slang#5994
Closes shader-slang#6437

Author: expipiplus1

source/slang/slang-parameter-binding.cpp

@@ -2383,6 +2383,21 @@ static RefPtr<TypeLayout> processEntryPointVaryingParameter(
     else if (auto declRefType = as<DeclRefType>(type))
     {
         auto declRef = declRefType->getDeclRef();
+
+        // 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.
+        const auto isExtern = declRef.getDecl()->hasModifier<ExternAttribute>() ||
+                              declRef.getDecl()->hasModifier<ExternModifier>();
+        if (isExtern)
+        {
+            if (const auto declDef = context->layoutContext.lookupExternDecl(declRef.getDecl()))
+            {
+                declRefType = declDef;
+                declRef = declRefType->getDeclRef();
+            }
+        }
+
         if (auto structDeclRef = declRef.as<StructDecl>())
         {
             RefPtr<StructTypeLayout> structLayout = new StructTypeLayout();

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>
@@ -5694,4 +5695,72 @@ GlobalGenericParamDecl* GenericParamTypeLayout::getGlobalGenericParamDecl()
     return rsDeclRef.getDecl();
 }
 
+DeclRefType* TypeLayoutContext::lookupExternDecl(Decl* decl)
+{
+    if (!externTypeMap)
+        buildExternTypeMap();
+    DeclRefType* ret = nullptr;
+    const auto mangledName = getMangledName(targetReq->getLinkage()->getASTBuilder(), decl);
+    externTypeMap->tryGetValue(mangledName, ret);
+    return ret;
+}
+
+void TypeLayoutContext::buildExternTypeMap()
+{
+    externTypeMap.emplace();
+    const auto linkage = targetReq->getLinkage();
+    auto astBuilder = linkage->getASTBuilder();
+
+    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

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* lookupExternDecl(Decl *decl);
+    void buildExternTypeMap();
+
     LayoutRulesImpl* getRules() { return rules; }
     LayoutRulesFamilyImpl* getRulesFamily() const { return rules->getLayoutRulesFamily(); }