Test

Author: jketema

swift/extractor/infra/SwiftBodyEmissionStrategy.cpp

@@ -2,6 +2,49 @@
 
 using namespace codeql;
 
+#include <swift/AST/ASTContext.h>
+#include <swift/AST/ClangModuleLoader.h>
+#include <clang/Basic/Module.h>
+
+namespace {
+const swift::ModuleDecl* getRealModuleOf(const swift::Decl* decl) {
+  auto* swiftModule = decl->getModuleContext();
+  auto* clangModule = swiftModule->findUnderlyingClangModule();
+
+  if (!clangModule) {
+    return swiftModule;
+  }
+
+  auto* clangModuleLoader = decl->getASTContext().getClangModuleLoader();
+
+  if (!clangModuleLoader) {
+    return swiftModule;
+  }
+
+  static std::unordered_map<const swift::Decl*, const swift::ModuleDecl*> cache;
+
+  if (auto result = cache.find(decl); result != cache.end()) {
+    return result->second;
+  }
+
+  for (const auto& submodule : clangModule->submodules()) {
+    if (auto* swiftSubmodule = clangModuleLoader->getWrapperForModule(submodule)) {
+      llvm::SmallVector<swift::Decl*> children;
+      swiftSubmodule->getTopLevelDecls(children);
+      for (const auto child : children) {
+        cache[child] = swiftSubmodule;
+      }
+    }
+  }
+
+  if (auto result = cache.find(decl); result != cache.end()) {
+    return result->second;
+  } else {
+    return swiftModule;
+  }
+}
+}  // namespace
+
 // In order to not emit duplicated entries for declarations, we restrict emission to only
 // Decls declared within the current "scope".
 // Depending on the whether we are extracting a primary source file or not the scope is defined as
@@ -12,7 +55,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();
+  auto module = getRealModuleOf(&decl);
   if (module != &currentModule) {
     return false;
   }