[SourceKit/InterfaceGen] Also print the cross-import overlays of a module in its interface.

include/swift/AST/Module.h

@@ -335,6 +335,31 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   void getDeclaredCrossImportBystanders(
       SmallVectorImpl<Identifier> &bystanderNames);
 
+  /// A lazily populated  mapping from each declared cross import overlay this
+  /// module transitively underlies to its bystander and immediate underlying
+  /// module.
+  llvm::SmallDenseMap<ModuleDecl *, std::pair<Identifier, ModuleDecl *>, 1>
+  declaredCrossImportsTransitive;
+
+  /// Determines if the given \p overlay is a declarared cross-import overlay of
+  /// this module, or an of its transitively declared overlay modules.
+  ///
+  /// This is used by tooling to map overlays to their underlying modules, and t
+  bool isUnderlyingModuleOfCrossImportOverlay(const ModuleDecl *overlay);
+
+  /// If \p overlay is a transitively declared cross-import overlay of this
+  /// module, gets the list of bystander modules that need to be imported
+  /// alongside this module for the overlay to be loaded.
+  void getAllBystandersForCrossImportOverlay(
+      ModuleDecl *overlay, SmallVectorImpl<Identifier> &bystanders);
+
+  /// Walks and loads the declared cross-import overlays of this module,
+  /// transitively, to find all overlays this module underlies.
+  ///
+  /// This is used by tooling to present these overlays as part of this module.
+  void findDeclaredCrossImportOverlaysTransitive(
+      SmallVectorImpl<ModuleDecl *> &overlays);
+
   /// Convenience accessor for clients that know what kind of file they're
   /// dealing with.
   SourceFile &getMainSourceFile(SourceFileKind expectedKind) const;

include/swift/AST/PrintOptions.h

@@ -432,6 +432,13 @@ struct PrintOptions {
   /// The information for converting archetypes to specialized types.
   llvm::Optional<TypeTransformContext> TransformContext;
 
+  /// Before printing the name of a ModuleDecl, this callback will be called and
+  /// the name of the ModuleDecl it returns will be printed instead. This is
+  /// currently used to present cross import overlays as if they were their
+  /// underlying module.
+  std::function<const ModuleDecl*(const ModuleDecl *)> mapModuleToUnderlying =
+    [] (const ModuleDecl *D) { return D; };
+
   bool PrintAsMember = false;
 
   /// Whether to print parameter specifiers as 'let' and 'var'.

lib/AST/ASTPrinter.cpp

@@ -3586,7 +3586,8 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
   template <typename T>
   void printModuleContext(T *Ty) {
     FileUnit *File = cast<FileUnit>(Ty->getDecl()->getModuleScopeContext());
-    ModuleDecl *Mod = File->getParentModule();
+    const ModuleDecl *Mod =
+        Options.mapModuleToUnderlying(File->getParentModule());
 
     Identifier Name = Mod->getName();
     if (Options.UseExportedModuleNames)

lib/AST/Module.cpp

@@ -1588,6 +1588,80 @@ void ModuleDecl::getDeclaredCrossImportBystanders(
     otherModules.push_back(std::get<0>(pair));
 }
 
+using TransitiveOverlays =
+    llvm::SmallDenseMap<ModuleDecl *, std::pair<Identifier, ModuleDecl *>, 1>;
+
+static void populateTransitiveCrossImports(ModuleDecl *base,
+                                           TransitiveOverlays &result) {
+  if (!result.empty() || !base->mightDeclareCrossImportOverlays())
+    return;
+
+  SmallVector<Identifier, 1> bystanders;
+  SmallVector<Identifier, 1> overlays;
+  SmallVector<ModuleDecl *, 1> worklist;
+  SourceLoc diagLoc; // ignored
+
+  worklist.push_back(base);
+  while (!worklist.empty()) {
+    ModuleDecl *current = worklist.back();
+    worklist.pop_back();
+    if (!current->mightDeclareCrossImportOverlays())
+      continue;
+    bystanders.clear();
+    current->getDeclaredCrossImportBystanders(bystanders);
+    for (Identifier bystander: bystanders) {
+      overlays.clear();
+      current->findDeclaredCrossImportOverlays(bystander, overlays, diagLoc);
+      for (Identifier overlay: overlays) {
+        if (!overlay.str().startswith("_"))
+          continue;
+        ModuleDecl *overlayMod =
+            base->getASTContext().getModuleByName(overlay.str());
+        if (!overlayMod)
+          continue;
+        if (result.insert({overlayMod, {bystander, current}}).second)
+          worklist.push_back(overlayMod);
+      }
+    }
+  }
+}
+
+bool ModuleDecl::isUnderlyingModuleOfCrossImportOverlay(
+    const ModuleDecl *overlay) {
+  if (!overlay->getNameStr().startswith("_"))
+    return false;
+
+  populateTransitiveCrossImports(this, declaredCrossImportsTransitive);
+  return declaredCrossImportsTransitive.find(overlay) !=
+      declaredCrossImportsTransitive.end();
+}
+
+void ModuleDecl::getAllBystandersForCrossImportOverlay(
+    ModuleDecl *overlay, SmallVectorImpl<Identifier> &bystanders) {
+  if (!overlay->getNameStr().startswith("_"))
+    return;
+
+  populateTransitiveCrossImports(this, declaredCrossImportsTransitive);
+
+  auto end = declaredCrossImportsTransitive.end();
+  for (auto i = declaredCrossImportsTransitive.find(overlay);
+       i != end;
+       i = declaredCrossImportsTransitive.find(i->second.second)) {
+    bystanders.push_back(i->second.first);
+  }
+}
+
+void ModuleDecl::findDeclaredCrossImportOverlaysTransitive(
+    SmallVectorImpl<ModuleDecl *> &overlayModules) {
+  populateTransitiveCrossImports(this, declaredCrossImportsTransitive);
+  std::transform(declaredCrossImportsTransitive.begin(),
+                 declaredCrossImportsTransitive.end(),
+                 std::back_inserter(overlayModules),
+                 [](TransitiveOverlays::iterator::value_type &i) {
+    return i.first;
+  });
+}
+
 namespace {
 struct OverlayFileContents {
   struct Module {