Swift: Work around assertion failures in mangler

swift/extractor/SwiftExtractor.cpp

@@ -170,7 +170,7 @@ static std::unordered_set<swift::ModuleDecl*> extractDeclarations(
                        bodyEmissionStrategy);
   auto topLevelDecls = getTopLevelDecls(module, primaryFile, lazyDeclaration);
   for (auto decl : topLevelDecls) {
-    if (decl->isUnavailable()) {
+    if (decl->isUnavailable() && !llvm::isa<swift::NominalTypeDecl>(decl)) {
       continue;
     }
     visitor.extract(decl);

swift/extractor/infra/SwiftBodyEmissionStrategy.cpp

@@ -12,8 +12,7 @@ using namespace codeql;
 //    same module one by one. In this mode, we restrict emission only to the same file ignoring
 //    all the other files.
 bool SwiftBodyEmissionStrategy::shouldEmitDeclBody(const swift::Decl& decl) {
-  auto module = decl.getModuleContext();
-  if (module != &currentModule) {
+  if (!currentTopLevelDecls.contains(&decl) && decl.getModuleContext() != &currentModule) {
     return false;
   }
   if (currentLazyDeclaration && currentLazyDeclaration != &decl) {

swift/extractor/infra/SwiftBodyEmissionStrategy.h

@@ -12,13 +12,19 @@ class SwiftBodyEmissionStrategy {
                             const swift::Decl* currentLazyDeclaration)
       : currentModule(currentModule),
         currentPrimarySourceFile(currentPrimarySourceFile),
-        currentLazyDeclaration(currentLazyDeclaration) {}
+        currentLazyDeclaration(currentLazyDeclaration) {
+    llvm::SmallVector<swift::Decl*> decls;
+    currentModule.getTopLevelDecls(decls);
+    currentTopLevelDecls.insert(decls.begin(), decls.end());
+  }
+
   bool shouldEmitDeclBody(const swift::Decl& decl);
 
  private:
   swift::ModuleDecl& currentModule;
   swift::SourceFile* currentPrimarySourceFile;
   const swift::Decl* currentLazyDeclaration;
+  std::unordered_set<const swift::Decl*> currentTopLevelDecls;
 };
 
 }  // namespace codeql

swift/extractor/mangler/SwiftMangler.cpp

@@ -8,6 +8,8 @@
 #include <swift/AST/ASTContext.h>
 #include <swift/AST/GenericEnvironment.h>
 #include <swift/AST/GenericParamList.h>
+#include <swift/AST/ClangModuleLoader.h>
+#include <clang/Basic/Module.h>
 
 using namespace codeql;
 
@@ -38,6 +40,8 @@ std::string_view getTypeKindStr(const swift::TypeBase* type) {
 
 }  // namespace
 
+std::unordered_map<const swift::Decl*, unsigned> SwiftMangler::preloadedExtensionIndexes;
+
 SwiftMangledName SwiftMangler::initMangled(const swift::TypeBase* type) {
   return {getTypeKindStr(type), '_'};
 }
@@ -109,32 +113,58 @@ unsigned SwiftMangler::getExtensionIndex(const swift::ExtensionDecl* decl,
   // indexes once for each encountered parent into the `preloadedExtensionIndexes` mapping.
   // Because we mangle declarations only once in a given trap/dispatcher context, we can safely
   // discard preloaded indexes on use
-  if (auto found = preloadedExtensionIndexes.extract(decl)) {
-    return found.mapped();
+  if (auto found = SwiftMangler::preloadedExtensionIndexes.find(decl);
+      found != SwiftMangler::preloadedExtensionIndexes.end()) {
+    return found->second;
   }
   if (auto parentModule = llvm::dyn_cast<swift::ModuleDecl>(parent)) {
     llvm::SmallVector<swift::Decl*> siblings;
     parentModule->getTopLevelDecls(siblings);
     indexExtensions(siblings);
+    if (auto clangModule = parentModule->findUnderlyingClangModule()) {
+      indexClangExtensions(clangModule, decl->getASTContext().getClangModuleLoader());
+    }
   } else if (auto iterableParent = llvm::dyn_cast<swift::IterableDeclContext>(parent)) {
     indexExtensions(iterableParent->getAllMembers());
   } else {
     // TODO use a generic logging handle for Swift entities here, once it's available
     CODEQL_ASSERT(false, "non-local context must be module or iterable decl context");
   }
-  auto found = preloadedExtensionIndexes.extract(decl);
+  auto found = SwiftMangler::preloadedExtensionIndexes.find(decl);
   // TODO use a generic logging handle for Swift entities here, once it's available
-  CODEQL_ASSERT(found, "extension not found within parent");
-  return found.mapped();
+  CODEQL_ASSERT(found != SwiftMangler::preloadedExtensionIndexes.end(),
+                "extension not found within parent");
+  return found->second;
 }
 
 void SwiftMangler::indexExtensions(llvm::ArrayRef<swift::Decl*> siblings) {
   auto index = 0u;
   for (auto sibling : siblings) {
     if (sibling->getKind() == swift::DeclKind::Extension) {
-      preloadedExtensionIndexes.emplace(sibling, index);
+      SwiftMangler::preloadedExtensionIndexes.emplace(sibling, index);
+      index += 2;
+    }
+  }
+}
+
+void SwiftMangler::indexClangExtensions(const clang::Module* clangModule,
+                                        swift::ClangModuleLoader* moduleLoader) {
+  if (!moduleLoader) {
+    return;
+  }
+
+  auto index = 1u;
+  for (const auto& submodule : clangModule->submodules()) {
+    if (auto* swiftSubmodule = moduleLoader->getWrapperForModule(submodule)) {
+      llvm::SmallVector<swift::Decl*> children;
+      swiftSubmodule->getTopLevelDecls(children);
+      for (const auto child : children) {
+        if (child->getKind() == swift::DeclKind::Extension) {
+          SwiftMangler::preloadedExtensionIndexes.emplace(child, index);
+          index += 2;
+        }
+      }
     }
-    ++index;
   }
 }
 

swift/extractor/mangler/SwiftMangler.h

@@ -106,13 +106,15 @@ class SwiftMangler : private swift::TypeVisitor<SwiftMangler, SwiftMangledName>,
   SwiftMangledName visitPackExpansionType(const swift::PackExpansionType* type);
 
  private:
-  std::unordered_map<const swift::Decl*, unsigned> preloadedExtensionIndexes;
+  static std::unordered_map<const swift::Decl*, unsigned> preloadedExtensionIndexes;
 
   virtual SwiftMangledName fetch(const swift::Decl* decl) = 0;
   virtual SwiftMangledName fetch(const swift::TypeBase* type) = 0;
   SwiftMangledName fetch(swift::Type type) { return fetch(type.getPointer()); }
 
   void indexExtensions(llvm::ArrayRef<swift::Decl*> siblings);
+  void indexClangExtensions(const clang::Module* clangModule,
+                            swift::ClangModuleLoader* moduleLoader);
   unsigned int getExtensionIndex(const swift::ExtensionDecl* decl, const swift::Decl* parent);
   static SwiftMangledName initMangled(const swift::TypeBase* type);
   SwiftMangledName initMangled(const swift::Decl* decl);