add support to getTopLevelDecls for clang submodules

QuietMisdreavus · QuietMisdreavus · commit 9d83f19aa558 · 2024-09-11T10:51:42.000-06:00
rdar://126031510
diff --git a/include/swift/AST/ClangNode.h b/include/swift/AST/ClangNode.h
@@ -96,6 +96,9 @@ class ClangNode {
   /// clang::ImportDecl or null if it's neither.
   const clang::Module *getClangModule() const;
 
+  /// Returns the owning clang module of this node, if it exists.
+  const clang::Module *getOwningClangModule() const;
+
   clang::SourceLocation getLocation() const;
   clang::SourceRange getSourceRange() const;
 
diff --git a/lib/AST/Decl.cpp b/lib/AST/Decl.cpp
@@ -128,6 +128,16 @@ const clang::Module *ClangNode::getClangModule() const {
   return nullptr;
 }
 
+const clang::Module *ClangNode::getOwningClangModule() const {
+  if (auto *M = getAsModule())
+    return M;
+  if (auto D = getAsDecl())
+    return D->getOwningModule();
+  if (auto MI = getAsModuleMacro())
+    return MI->getOwningModule();
+  return nullptr;
+}
+
 void ClangNode::dump() const {
   if (auto D = getAsDecl())
     D->dump();
diff --git a/lib/ClangImporter/ClangImporter.cpp b/lib/ClangImporter/ClangImporter.cpp
@@ -3050,6 +3050,15 @@ static bool isDeclaredInModule(const ClangModuleUnit *ModuleFilter,
   if (VD->getModuleContext()->getName().str() == CLANG_HEADER_MODULE_NAME) {
     return true;
   }
+  // Because the ClangModuleUnit saved as a decl context will be saved as the top-level module, but
+  // the ModuleFilter we're given might be a submodule (if a submodule was passed to
+  // getTopLevelDecls, for example), we should compare the underlying Clang modules to determine
+  // module membership.
+  if (auto ClangNode = VD->getClangNode()) {
+    if (auto *ClangModule = ClangNode.getOwningClangModule()) {
+      return ModuleFilter->getClangModule() == ClangModule;
+    }
+  }
   auto ContainingUnit = VD->getDeclContext()->getModuleScopeContext();
   return ModuleFilter == ContainingUnit;
 }
@@ -3222,7 +3231,7 @@ class FilteringDeclaredDeclConsumer : public swift::VisibleDeclConsumer {
   FilteringDeclaredDeclConsumer(swift::VisibleDeclConsumer &consumer,
                                 const ClangModuleUnit *CMU)
       : NextConsumer(consumer), ModuleFilter(CMU) {
-    assert(CMU && CMU->isTopLevel() && "Only top-level modules supported");
+    assert(CMU);
   }
 
   void foundDecl(ValueDecl *VD, DeclVisibilityKind Reason,
@@ -3596,10 +3605,12 @@ void ClangModuleUnit::getTopLevelDecls(SmallVectorImpl<Decl*> &results) const {
 
     // Add the extensions produced by importing categories.
     for (auto category : lookupTable->categories()) {
-      if (auto extension = cast_or_null<ExtensionDecl>(
-              owner.importDecl(category, owner.CurrentVersion,
-                               /*UseCanonical*/false))) {
-        results.push_back(extension);
+      if (category->getOwningModule() == clangModule) {
+        if (auto extension = cast_or_null<ExtensionDecl>(
+          owner.importDecl(category, owner.CurrentVersion,
+                          /*UseCanonical*/false))) {
+          results.push_back(extension);
+        }
       }
     }
 
diff --git a/lib/IDE/ModuleInterfacePrinting.cpp b/lib/IDE/ModuleInterfacePrinting.cpp
@@ -446,9 +446,6 @@ void swift::ide::printModuleInterface(
        const PrintOptions &Options,
        const bool PrintSynthesizedExtensions) {
 
-  // Clang submodules aren't handled well by `getDisplayDecls()` (no decls are
-  // returned), so map them to their top-level module and filter out the extra
-  // results below.
   const clang::Module *TargetClangMod = TargetMod->findUnderlyingClangModule();
   ModuleDecl *TopLevelMod = TargetMod->getTopLevelModule();
   bool IsSubmodule = TargetMod != TopLevelMod;
@@ -460,8 +457,8 @@ void swift::ide::printModuleInterface(
   auto AdjustedOptions = Options;
   adjustPrintOptions(AdjustedOptions);
 
-  SmallVector<Decl *, 1> Decls;
-  swift::getTopLevelDeclsForDisplay(TopLevelMod, Decls);
+  SmallVector<ModuleDecl *, 1> ModuleList;
+  ModuleList.push_back(TargetMod);
 
   SmallVector<ImportDecl *, 1> ImportDecls;
   llvm::DenseSet<const clang::Module *> ClangModulesForImports;
@@ -485,6 +482,10 @@ void swift::ide::printModuleInterface(
 
         ClangDecls.insert({ CM, {} });
 
+        if (CM != TargetClangMod)
+          if (auto *OwningModule = Importer.getWrapperForModule(CM))
+            ModuleList.push_back(OwningModule);
+
         // If we're supposed to visit submodules, add them now.
         if (TraversalOptions & ModuleTraversal::VisitSubmodules) {
           for (clang::Module * submodule: CM->submodules()) {
@@ -499,6 +500,12 @@ void swift::ide::printModuleInterface(
     }
   }
 
+  SmallVector<Decl *, 1> Decls;
+
+  for (ModuleDecl *M : ModuleList) {
+    swift::getTopLevelDeclsForDisplay(M, Decls);
+  }
+
   // Collect those submodules that are actually imported but have no import
   // decls in the module.
   llvm::SmallPtrSet<const clang::Module *, 16> NoImportSubModules;
diff --git a/lib/SymbolGraphGen/SymbolGraph.cpp b/lib/SymbolGraphGen/SymbolGraph.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/AST/DeclObjC.h"
+#include "clang/Basic/Module.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Comment.h"
 #include "swift/AST/Decl.h"
@@ -805,7 +806,15 @@ bool SymbolGraph::isUnconditionallyUnavailableOnAllPlatforms(const Decl *D) cons
 bool SymbolGraph::canIncludeDeclAsNode(const Decl *D) const {
   // If this decl isn't in this module or module that this module imported with `@_exported`, don't record it,
   // as it will appear elsewhere in its module's symbol graph.
-  if (D->getModuleContext()->getName() != M.getName() && !Walker.isConsideredExportedImported(D)) {
+
+  // If a Clang decl was declared in a submodule, the Swift decl's context will still point to the
+  // top-level module. Instead, we need to probe the owning module on the Clang side, which will
+  // correctly point to the submodule.
+  auto RealModuleName = (std::string)D->getModuleContext()->getName();
+  if (auto *ClangDecl = D->getClangDecl())
+    if (auto *ClangModule = ClangDecl->getOwningModule())
+      RealModuleName = ClangModule->Name;
+  if (RealModuleName != (std::string)M.getName() && !Walker.isConsideredExportedImported(D)) {
     return false;
   }
 
diff --git a/lib/SymbolGraphGen/SymbolGraphASTWalker.cpp b/lib/SymbolGraphGen/SymbolGraphASTWalker.cpp
@@ -10,7 +10,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "clang/AST/DeclObjC.h"
+#include "clang/Basic/Module.h"
 #include "llvm/ADT/StringSwitch.h"
+#include "swift/AST/ASTContext.h"
+#include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Module.h"
 #include "swift/AST/ProtocolConformance.h"
@@ -57,20 +61,32 @@ SymbolGraphASTWalker::SymbolGraphASTWalker(
       QualifiedExportedImports(QualifiedExportedImports),
       MainGraph(*this, M, std::nullopt, Ctx) {}
 
+ModuleDecl *SymbolGraphASTWalker::getRealModuleOf(const Decl *D) const {
+  ModuleDecl *Module = D->getModuleContext();
+  if (auto *ClangDecl = D->getClangDecl())
+    if (auto *ClangModule = ClangDecl->getOwningModule())
+      if (auto *ClangModuleLoader = D->getASTContext().getClangModuleLoader())
+        if (auto *M = ClangModuleLoader->getWrapperForModule(ClangModule))
+          Module = M;
+
+  return Module;
+}
+
 /// Get a "sub" symbol graph for the parent module of a type that
 /// the main module `M` is extending.
 SymbolGraph *SymbolGraphASTWalker::getModuleSymbolGraph(const Decl *D) {
-  auto *M = D->getModuleContext();
+  auto *M = getRealModuleOf(D);
   const auto *DC = D->getDeclContext();
   SmallVector<const NominalTypeDecl *, 2> ParentTypes = {};
   const Decl *ExtendedNominal = nullptr;
   while (DC) {
-    M = DC->getParentModule();
     if (const auto *NTD = dyn_cast_or_null<NominalTypeDecl>(DC->getAsDecl())) {
       DC = NTD->getDeclContext();
+      M = getRealModuleOf(NTD);
       ParentTypes.push_back(NTD);
     } else if (const auto *Ext = dyn_cast_or_null<ExtensionDecl>(DC->getAsDecl())) {
       DC = Ext->getExtendedNominal()->getDeclContext();
+      M = getRealModuleOf(Ext->getExtendedNominal());
       if (!ExtendedNominal)
         ExtendedNominal = Ext->getExtendedNominal();
     } else {
diff --git a/lib/SymbolGraphGen/SymbolGraphASTWalker.h b/lib/SymbolGraphGen/SymbolGraphASTWalker.h
@@ -130,6 +130,10 @@ struct SymbolGraphASTWalker : public SourceEntityWalker {
   /// extension block symbol, or if its members should be directly associated
   /// with its extended nominal.
   virtual bool shouldBeRecordedAsExtension(const ExtensionDecl *ED) const;
+
+  /// Returns the owning module of the given decl. Loads the module from Clang if necessary, to
+  /// correctly fetch owning submodules.
+  virtual ModuleDecl *getRealModuleOf(const Decl *D) const;
 };
 
 } // end namespace symbolgraphgen
diff --git a/test/SourceKit/DocSupport/doc_clang_module.swift.sub.response b/test/SourceKit/DocSupport/doc_clang_module.swift.sub.response
diff --git a/test/SourceKit/InterfaceGen/gen_clang_module.swift.sub.response b/test/SourceKit/InterfaceGen/gen_clang_module.swift.sub.response
diff --git a/test/SymbolGraph/ClangImporter/Submodules.swift b/test/SymbolGraph/ClangImporter/Submodules.swift
