🍒[6.0][SymbolGraphGen] Correctly handle exported imports in swift-symbolgraph-extract

include/swift/AST/Module.h

@@ -18,6 +18,7 @@
 #define SWIFT_MODULE_H
 
 #include "swift/AST/AccessNotes.h"
+#include "swift/AST/AttrKind.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/DeclContext.h"
 #include "swift/AST/Identifier.h"
@@ -551,7 +552,7 @@ class ModuleDecl
   ModuleDecl *getUnderlyingModuleIfOverlay() const;
 
   /// Returns true if this module is the Clang overlay of \p other.
-  bool isClangOverlayOf(ModuleDecl *other);
+  bool isClangOverlayOf(ModuleDecl *other) const;
 
   /// Returns true if this module is the same module or either module is a clang
   /// overlay of the other.
@@ -1084,6 +1085,24 @@ class ModuleDecl
   /// for that.
   void getDisplayDecls(SmallVectorImpl<Decl*> &results, bool recursive = false) const;
 
+  struct ImportCollector {
+    SmallPtrSet<const ModuleDecl *, 4> imports;
+    llvm::SmallDenseMap<const ModuleDecl *, SmallPtrSet<Decl *, 4>, 4>
+        qualifiedImports;
+    AccessLevel minimumDocVisibility = AccessLevel::Private;
+    llvm::function_ref<bool(const ModuleDecl *)> importFilter = nullptr;
+
+    void collect(const ImportedModule &importedModule);
+
+    ImportCollector() = default;
+    ImportCollector(AccessLevel minimumDocVisibility)
+        : minimumDocVisibility(minimumDocVisibility) {}
+  };
+
+  void
+  getDisplayDeclsRecursivelyAndImports(SmallVectorImpl<Decl *> &results,
+                                       ImportCollector &importCollector) const;
+
   using LinkLibraryCallback = llvm::function_ref<void(LinkLibrary)>;
 
   /// Generate the list of libraries needed to link this module, based on its
@@ -1263,12 +1282,6 @@ inline SourceLoc extractNearestSourceLoc(const ModuleDecl *mod) {
   return extractNearestSourceLoc(static_cast<const Decl *>(mod));
 }
 
-/// Collects modules that this module imports via `@_exported import`.
-void collectParsedExportedImports(const ModuleDecl *M,
-                                  SmallPtrSetImpl<ModuleDecl *> &Imports,
-                                  llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> &QualifiedImports,
-                                  llvm::function_ref<bool(AttributedImport<ImportedModule>)> includeImport = nullptr);
-
 /// If the import that would make the given declaration visibile is absent,
 /// emit a diagnostic and a fix-it suggesting adding the missing import.
 bool diagnoseMissingImportForMember(const ValueDecl *decl,

include/swift/Option/Options.td

@@ -1573,6 +1573,10 @@ def output_dir : Separate<["-"], "output-dir">,
   HelpText<"Output directory">,
   MetaVarName<"<dir>">;
 
+def experimental_allowed_reexported_modules: CommaJoined<["-"], "experimental-allowed-reexported-modules=">,
+  Flags<[NoDriverOption, SwiftSymbolGraphExtractOption]>,
+  HelpText<"Allow reexporting symbols from the provided modules if they are themselves exported from the main module. This is a comma separated list of module names.">;
+
 def skip_synthesized_members: Flag<[ "-" ], "skip-synthesized-members">,
   Flags<[NoDriverOption, SwiftSymbolGraphExtractOption]>,
   HelpText<"Skip members inherited through classes or default implementations">;

include/swift/Sema/IDETypeChecking.h

@@ -162,6 +162,15 @@ namespace swift {
   void
   getTopLevelDeclsForDisplay(ModuleDecl *M, SmallVectorImpl<Decl*> &Results, bool Recursive = false);
 
+  /// Get all of the top-level declarations that should be printed as part of
+  /// this module. This may force synthesis of top-level declarations that
+  /// \p getDisplayDeclsForModule would only return if previous
+  /// work happened to have synthesized them.
+  void getTopLevelDeclsForDisplay(
+      ModuleDecl *M, SmallVectorImpl<Decl *> &Results,
+      llvm::function_ref<void(ModuleDecl *, SmallVectorImpl<Decl *> &)>
+          getDisplayDeclsForModule);
+
   struct ExtensionInfo {
     // The extension with the declarations to apply.
     ExtensionDecl *Ext;

include/swift/SymbolGraphGen/SymbolGraphOptions.h

@@ -11,6 +11,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/TargetParser/Triple.h"
+#include "llvm/ADT/ArrayRef.h"
+
 #include "swift/AST/AttrKind.h"
 
 #ifndef SWIFT_SYMBOLGRAPHGEN_SYMBOLGRAPHOPTIONS_H
@@ -63,6 +65,10 @@ struct SymbolGraphOptions {
   /// but SourceKit should be able to load the information when pulling symbol
   /// information for individual queries.
   bool PrintPrivateStdlibSymbols = false;
+
+  /// If this has a value specifies an explicit allow list of reexported module
+  /// names that should be included symbol graph.
+  std::optional<llvm::ArrayRef<StringRef>> AllowedReexportedModules = {};
 };
 
 } // end namespace symbolgraphgen

lib/AST/Module.cpp

@@ -26,6 +26,7 @@
 #include "swift/AST/ExistentialLayout.h"
 #include "swift/AST/FileUnit.h"
 #include "swift/AST/GenericEnvironment.h"
+#include "swift/AST/Import.h"
 #include "swift/AST/ImportCache.h"
 #include "swift/AST/LazyResolver.h"
 #include "swift/AST/LinkLibrary.h"
@@ -39,6 +40,7 @@
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/SynthesizedFileUnit.h"
+#include "swift/AST/Type.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/Basic/Compiler.h"
 #include "swift/Basic/SourceManager.h"
@@ -1287,42 +1289,6 @@ bool ModuleDecl::shouldCollectDisplayDecls() const {
   return true;
 }
 
-void swift::collectParsedExportedImports(const ModuleDecl *M,
-                                         SmallPtrSetImpl<ModuleDecl *> &Imports,
-                                         llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> &QualifiedImports,
-                                         llvm::function_ref<bool(AttributedImport<ImportedModule>)> includeImport) {
-  for (const FileUnit *file : M->getFiles()) {
-    if (const SourceFile *source = dyn_cast<SourceFile>(file)) {
-      if (source->hasImports()) {
-        for (auto import : source->getImports()) {
-          if (import.options.contains(ImportFlags::Exported) &&
-              (!includeImport || includeImport(import)) &&
-              import.module.importedModule->shouldCollectDisplayDecls()) {
-            auto *TheModule = import.module.importedModule;
-
-            if (import.module.getAccessPath().size() > 0) {
-              if (QualifiedImports.find(TheModule) == QualifiedImports.end()) {
-                QualifiedImports.try_emplace(TheModule);
-              }
-              auto collectDecls = [&](ValueDecl *VD,
-                                      DeclVisibilityKind reason) {
-                if (reason == DeclVisibilityKind::VisibleAtTopLevel)
-                  QualifiedImports[TheModule].insert(VD);
-              };
-              auto consumer = makeDeclConsumer(std::move(collectDecls));
-              TheModule->lookupVisibleDecls(
-                  import.module.getAccessPath(), consumer,
-                  NLKind::UnqualifiedLookup);
-            } else if (!Imports.contains(TheModule)) {
-              Imports.insert(TheModule);
-            }
-          }
-        }
-      }
-    }
-  }
-}
-
 void ModuleDecl::getLocalTypeDecls(SmallVectorImpl<TypeDecl*> &Results) const {
   FORWARD(getLocalTypeDecls, (Results));
 }
@@ -1543,40 +1509,101 @@ SourceFile::getExternalRawLocsForDecl(const Decl *D) const {
   return Result;
 }
 
-void ModuleDecl::getDisplayDecls(SmallVectorImpl<Decl*> &Results, bool Recursive) const {
-  if (Recursive && isParsedModule(this)) {
-    SmallPtrSet<ModuleDecl *, 4> Modules;
-    llvm::SmallDenseMap<ModuleDecl *, SmallPtrSet<Decl *, 4>, 4> QualifiedImports;
-    collectParsedExportedImports(this, Modules, QualifiedImports);
-    for (const auto &QI : QualifiedImports) {
-      auto Module = QI.getFirst();
-      if (Modules.contains(Module)) continue;
+void ModuleDecl::ImportCollector::collect(
+    const ImportedModule &importedModule) {
+  auto *module = importedModule.importedModule;
 
-      auto &Decls = QI.getSecond();
-      Results.append(Decls.begin(), Decls.end());
-    }
-    for (const ModuleDecl *import : Modules) {
-      import->getDisplayDecls(Results, Recursive);
+  if (!module->shouldCollectDisplayDecls())
+    return;
+
+  if (importFilter && !importFilter(module))
+    return;
+
+  if (importedModule.getAccessPath().size() > 0) {
+    auto collectDecls = [&](ValueDecl *VD, DeclVisibilityKind reason) {
+      if (reason == DeclVisibilityKind::VisibleAtTopLevel)
+        this->qualifiedImports[module].insert(VD);
+    };
+    auto consumer = makeDeclConsumer(std::move(collectDecls));
+    module->lookupVisibleDecls(importedModule.getAccessPath(), consumer,
+                               NLKind::UnqualifiedLookup);
+  } else {
+    imports.insert(module);
+  }
+}
+
+static void
+collectExportedImports(const ModuleDecl *module,
+                       ModuleDecl::ImportCollector &importCollector) {
+  for (const FileUnit *file : module->getFiles()) {
+    if (const SourceFile *source = dyn_cast<SourceFile>(file)) {
+      if (source->hasImports()) {
+        for (const auto &import : source->getImports()) {
+          if (import.options.contains(ImportFlags::Exported) &&
+              import.docVisibility.value_or(AccessLevel::Public) >=
+                  importCollector.minimumDocVisibility) {
+            importCollector.collect(import.module);
+            collectExportedImports(import.module.importedModule,
+                                   importCollector);
+          }
+        }
+      }
+    } else {
+      SmallVector<ImportedModule, 8> exportedImports;
+      file->getImportedModules(exportedImports,
+                               ModuleDecl::ImportFilterKind::Exported);
+      for (const auto &im : exportedImports) {
+        // Skip collecting the underlying clang module as we already have the relevant import.
+        if (module->isClangOverlayOf(im.importedModule))
+          continue;
+        importCollector.collect(im);
+        collectExportedImports(im.importedModule, importCollector);
+      }
     }
   }
-  // FIXME: Should this do extra access control filtering?
-  FORWARD(getDisplayDecls, (Results));
+}
 
-#ifndef NDEBUG
+void ModuleDecl::getDisplayDecls(SmallVectorImpl<Decl*> &Results, bool Recursive) const {
   if (Recursive) {
-    llvm::DenseSet<Decl *> visited;
-    for (auto *D : Results) {
-      // decls synthesized from implicit clang decls may appear multiple times;
-      // e.g. if multiple modules with underlying clang modules are re-exported.
-      // including duplicates of these is harmless, so skip them when counting
-      // this assertion
-      if (const auto *CD = D->getClangDecl()) {
-        if (CD->isImplicit()) continue;
-      }
+    ImportCollector importCollector;
+    this->getDisplayDeclsRecursivelyAndImports(Results, importCollector);
+  } else {
+    // FIXME: Should this do extra access control filtering?
+    FORWARD(getDisplayDecls, (Results));
+  }
+}
+
+void ModuleDecl::getDisplayDeclsRecursivelyAndImports(
+    SmallVectorImpl<Decl *> &results, ImportCollector &importCollector) const {
+  this->getDisplayDecls(results, /*Recursive=*/false);
+
+  // Look up imports recursively.
+  collectExportedImports(this, importCollector);
+  for (const auto &QI : importCollector.qualifiedImports) {
+    auto Module = QI.getFirst();
+    if (importCollector.imports.contains(Module))
+      continue;
 
-      auto inserted = visited.insert(D).second;
-      assert(inserted && "there should be no duplicate decls");
+    auto &Decls = QI.getSecond();
+    results.append(Decls.begin(), Decls.end());
+  }
+
+  for (const ModuleDecl *import : importCollector.imports)
+    import->getDisplayDecls(results);
+
+#ifndef NDEBUG
+  llvm::DenseSet<Decl *> visited;
+  for (auto *D : results) {
+    // decls synthesized from implicit clang decls may appear multiple times;
+    // e.g. if multiple modules with underlying clang modules are re-exported.
+    // including duplicates of these is harmless, so skip them when counting
+    // this assertion
+    if (const auto *CD = D->getClangDecl()) {
+      if (CD->isImplicit())
+        continue;
     }
+    auto inserted = visited.insert(D).second;
+    assert(inserted && "there should be no duplicate decls");
   }
 #endif
 }
@@ -2393,7 +2420,7 @@ ModuleDecl::getDeclaringModuleAndBystander() {
   return *(declaringModuleAndBystander = {nullptr, Identifier()});
 }
 
-bool ModuleDecl::isClangOverlayOf(ModuleDecl *potentialUnderlying) {
+bool ModuleDecl::isClangOverlayOf(ModuleDecl *potentialUnderlying) const {
   return getUnderlyingModuleIfOverlay() == potentialUnderlying;
 }
 

lib/DriverTool/swift_symbolgraph_extract_main.cpp

@@ -163,6 +163,13 @@ int swift_symbolgraph_extract_main(ArrayRef<const char *> Args,
     }
   }
 
+  SmallVector<StringRef, 4> AllowedRexports;
+  if (auto *A =
+          ParsedArgs.getLastArg(OPT_experimental_allowed_reexported_modules)) {
+    for (const auto *val : A->getValues())
+      AllowedRexports.emplace_back(val);
+  }
+
   symbolgraphgen::SymbolGraphOptions Options;
   Options.OutputDir = OutputDir;
   Options.Target = Target;
@@ -175,6 +182,7 @@ int swift_symbolgraph_extract_main(ArrayRef<const char *> Args,
   Options.EmitExtensionBlockSymbols =
       ParsedArgs.hasFlag(OPT_emit_extension_block_symbols,
                          OPT_omit_extension_block_symbols, /*default=*/false);
+  Options.AllowedReexportedModules = AllowedRexports;
 
   if (auto *A = ParsedArgs.getLastArg(OPT_minimum_access_level)) {
     Options.MinimumAccessLevel =

lib/IDE/IDETypeChecking.cpp

@@ -10,9 +10,10 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/Sema/IDETypeChecking.h"
+#include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTDemangler.h"
 #include "swift/AST/ASTPrinter.h"
-#include "swift/AST/ASTContext.h"
 #include "swift/AST/Attr.h"
 #include "swift/AST/Decl.h"
 #include "swift/AST/Expr.h"
@@ -25,20 +26,30 @@
 #include "swift/AST/Requirement.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/Types.h"
-#include "swift/Sema/IDETypeChecking.h"
-#include "swift/Sema/IDETypeCheckingRequests.h"
-#include "swift/IDE/SourceEntityWalker.h"
 #include "swift/IDE/IDERequests.h"
+#include "swift/IDE/SourceEntityWalker.h"
 #include "swift/Parse/Lexer.h"
+#include "swift/Sema/IDETypeCheckingRequests.h"
+#include "llvm/ADT/SmallVector.h"
 
 using namespace swift;
 
-void
-swift::getTopLevelDeclsForDisplay(ModuleDecl *M,
-                                  SmallVectorImpl<Decl*> &Results,
-                                  bool Recursive) {
+void swift::getTopLevelDeclsForDisplay(ModuleDecl *M,
+                                       SmallVectorImpl<Decl *> &Results,
+                                       bool Recursive) {
+  auto getDisplayDeclsForModule =
+      [Recursive](ModuleDecl *M, SmallVectorImpl<Decl *> &Results) {
+        M->getDisplayDecls(Results, Recursive);
+      };
+  getTopLevelDeclsForDisplay(M, Results, std::move(getDisplayDeclsForModule));
+}
+
+void swift::getTopLevelDeclsForDisplay(
+    ModuleDecl *M, SmallVectorImpl<Decl *> &Results,
+    llvm::function_ref<void(ModuleDecl *, SmallVectorImpl<Decl *> &)>
+        getDisplayDeclsForModule) {
   auto startingSize = Results.size();
-  M->getDisplayDecls(Results, Recursive);
+  getDisplayDeclsForModule(M, Results);
 
   // Force Sendable on all public types, which might synthesize some extensions.
   // FIXME: We can remove this if @_nonSendable stops creating extensions.
@@ -48,20 +59,20 @@ swift::getTopLevelDeclsForDisplay(ModuleDecl *M,
       // Restrict this logic to public and package types. Non-public types
       // may refer to implementation details and fail at deserialization.
       auto accessScope = NTD->getFormalAccessScope();
-      if (!M->isMainModule() &&
-          !accessScope.isPublic() && !accessScope.isPackage())
+      if (!M->isMainModule() && !accessScope.isPublic() &&
+          !accessScope.isPackage())
         continue;
 
       auto proto = M->getASTContext().getProtocol(KnownProtocolKind::Sendable);
       if (proto)
-        (void) M->lookupConformance(NTD->getDeclaredInterfaceType(), proto);
+        (void)M->lookupConformance(NTD->getDeclaredInterfaceType(), proto);
     }
   }
 
   // Remove what we fetched and fetch again, possibly now with additional
   // extensions.
   Results.resize(startingSize);
-  M->getDisplayDecls(Results, Recursive);
+  getDisplayDeclsForModule(M, Results);
 }
 
 static bool shouldPrintAsFavorable(const Decl *D, const PrintOptions &Options) {