Merge pull request #30758 from nathawes/cross-import-fixes

Author: swift-ci

include/swift/AST/Module.h

@@ -338,26 +338,41 @@ 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;
+private:
+  /// A cache of this module's underlying module and required bystander if it's
+  /// an underscored cross-import overlay.
+  Optional<std::pair<ModuleDecl *, Identifier>> declaringModuleAndBystander;
 
-  /// 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 this module is an underscored cross import overlay, gets the underlying
+  /// module that declared it (which may itself be a cross-import overlay),
+  /// along with the name of the required bystander module. Used by tooling to
+  /// present overlays as if they were part of their underlying module.
+  std::pair<ModuleDecl *, Identifier> getDeclaringModuleAndBystander();
+
+public:
+
+  /// Returns true if this module is an underscored cross import overlay
+  /// declared by \p other, either directly or transitively (via intermediate
+  /// cross-import overlays - for cross-imports involving more than two
+  /// modules).
+  bool isCrossImportOverlayOf(ModuleDecl *other);
+
+  /// If this module is an underscored cross-import overlay, returns the
+  /// non-underscored underlying module that declares it as an overlay, either
+  /// directly or transitively (via intermediate cross-import overlays - for
+  /// cross-imports involving more than two modules).
+  ModuleDecl *getDeclaringModuleIfCrossImportOverlay();
+
+  /// If this module is an underscored cross-import overlay of \p declaring
+  /// either directly or transitively, populates \p bystanderNames with the set
+  /// of bystander modules that must be present alongside \p declaring for
+  /// the overlay to be imported and returns true. Returns false otherwise.
+  bool getRequiredBystandersIfCrossImportOverlay(
+      ModuleDecl *declaring, SmallVectorImpl<Identifier> &bystanderNames);
 
-  /// 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.
+  /// Walks and loads the declared, underscored 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(

include/swift/AST/PrintOptions.h

@@ -432,12 +432,10 @@ 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; };
+  /// Whether cross-import overlay modules are printed with their own name (e.g.
+  /// _MyFrameworkYourFrameworkAdditions) or that of their underlying module
+  /// (e.g.  MyFramework).
+  bool MapCrossImportOverlaysToDeclaringModule = false;
 
   bool PrintAsMember = false;
 
@@ -525,6 +523,7 @@ struct PrintOptions {
     result.PrintDocumentationComments = true;
     result.SkipUnderscoredKeywords = true;
     result.EnumRawValues = EnumRawValueMode::PrintObjCOnly;
+    result.MapCrossImportOverlaysToDeclaringModule = true;
     return result;
   }
 

include/swift/AST/SourceFile.h

@@ -194,16 +194,6 @@ class SourceFile final : public FileUnit {
   /// mechanism which is not SourceFile-dependent.)
   SeparatelyImportedOverlayMap separatelyImportedOverlays;
 
-  using SeparatelyImportedOverlayReverseMap =
-    llvm::SmallDenseMap<ModuleDecl *, ModuleDecl *>;
-
-  /// A lazily populated mapping from a separately imported overlay to its
-  /// underlying shadowed module.
-  ///
-  /// This is used by tooling to substitute the name of the underlying module
-  /// wherever the overlay's name would otherwise be reported.
-  SeparatelyImportedOverlayReverseMap separatelyImportedOverlaysReversed;
-
   /// A pointer to PersistentParserState with a function reference to its
   /// deleter to handle the fact that it's forward declared.
   using ParserStatePtr =
@@ -384,7 +374,6 @@ class SourceFile final : public FileUnit {
   /// \returns true if the overlay was added; false if it already existed.
   bool addSeparatelyImportedOverlay(ModuleDecl *overlay,
                                     ModuleDecl *declaring) {
-    separatelyImportedOverlaysReversed.clear();
     return std::get<1>(separatelyImportedOverlays[declaring].insert(overlay));
   }
 
@@ -402,12 +391,6 @@ class SourceFile final : public FileUnit {
     overlays.append(value.begin(), value.end());
   }
 
-  /// Retrieves a module shadowed by the provided separately imported overlay
-  /// \p shadowed. If such a module is returned, it should be presented to users
-  /// as owning the symbols in \p overlay.
-  ModuleDecl *
-  getModuleShadowedBySeparatelyImportedOverlay(const ModuleDecl *overlay);
-
   void cacheVisibleDecls(SmallVectorImpl<ValueDecl *> &&globals) const;
   const SmallVectorImpl<ValueDecl *> &getCachedVisibleDecls() const;
 

lib/AST/ASTPrinter.cpp

@@ -2153,7 +2153,16 @@ void PrintAST::visitImportDecl(ImportDecl *decl) {
   interleave(decl->getFullAccessPath(),
              [&](const ImportDecl::AccessPathElement &Elem) {
                if (!Mods.empty()) {
-                 Printer.printModuleRef(Mods.front(), Elem.Item);
+                 Identifier Name = Elem.Item;
+                 if (Options.MapCrossImportOverlaysToDeclaringModule) {
+                   if (auto *MD = Mods.front().getAsSwiftModule()) {
+                     ModuleDecl *Declaring = const_cast<ModuleDecl*>(MD)
+                       ->getDeclaringModuleIfCrossImportOverlay();
+                     if (Declaring)
+                       Name = Declaring->getName();
+                   }
+                 }
+                 Printer.printModuleRef(Mods.front(), Name);
                  Mods = Mods.slice(1);
                } else {
                  Printer << Elem.Item.str();
@@ -3591,8 +3600,12 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
   template <typename T>
   void printModuleContext(T *Ty) {
     FileUnit *File = cast<FileUnit>(Ty->getDecl()->getModuleScopeContext());
-    const ModuleDecl *Mod =
-        Options.mapModuleToUnderlying(File->getParentModule());
+    ModuleDecl *Mod = File->getParentModule();
+
+    if (Options.MapCrossImportOverlaysToDeclaringModule) {
+      if (ModuleDecl *Declaring = Mod->getDeclaringModuleIfCrossImportOverlay())
+        Mod = Declaring;
+    }
 
     Identifier Name = Mod->getName();
     if (Options.UseExportedModuleNames)

lib/AST/Module.cpp

@@ -1776,78 +1776,118 @@ 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;
+void ModuleDecl::findDeclaredCrossImportOverlaysTransitive(
+    SmallVectorImpl<ModuleDecl *> &overlayModules) {
   SmallVector<ModuleDecl *, 1> worklist;
-  SourceLoc diagLoc; // ignored
+  SmallPtrSet<ModuleDecl *, 1> seen;
+  SourceLoc unused;
 
-  worklist.push_back(base);
+  worklist.push_back(this);
   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);
+    for (auto &pair: current->declaredCrossImports) {
+      Identifier &bystander = std::get<0>(pair);
+      for (auto *file: std::get<1>(pair)) {
+        auto overlays = file->getOverlayModuleNames(current, unused, bystander);
+        for (Identifier overlay: overlays) {
+          // We don't present non-underscored overlays as part of the underlying
+          // module, so ignore them.
+          if (!overlay.str().startswith("_"))
+            continue;
+          ModuleDecl *overlayMod =
+              getASTContext().getModuleByName(overlay.str());
+          if (!overlayMod)
+            continue;
+          if (seen.insert(overlayMod).second) {
+            overlayModules.push_back(overlayMod);
+            worklist.push_back(overlayMod);
+          }
+        }
       }
     }
   }
 }
 
-bool ModuleDecl::isUnderlyingModuleOfCrossImportOverlay(
-    const ModuleDecl *overlay) {
-  if (!overlay->getNameStr().startswith("_"))
-    return false;
+std::pair<ModuleDecl *, Identifier>
+ModuleDecl::getDeclaringModuleAndBystander() {
+  if (declaringModuleAndBystander)
+    return *declaringModuleAndBystander;
 
-  populateTransitiveCrossImports(this, declaredCrossImportsTransitive);
-  return declaredCrossImportsTransitive.find(overlay) !=
-      declaredCrossImportsTransitive.end();
-}
+  if (!hasUnderscoredNaming())
+    return *(declaringModuleAndBystander = {nullptr, Identifier()});
 
-void ModuleDecl::getAllBystandersForCrossImportOverlay(
-    ModuleDecl *overlay, SmallVectorImpl<Identifier> &bystanders) {
-  if (!overlay->getNameStr().startswith("_"))
-    return;
+  // Search the transitive set of imported @_exported modules to see if any have
+  // this module as their overlay.
+  SmallPtrSet<ModuleDecl *, 16> seen;
+  SmallVector<ModuleDecl::ImportedModule, 16> imported;
+  SmallVector<ModuleDecl::ImportedModule, 16> furtherImported;
+  ModuleDecl *overlayModule = this;
+
+  getImportedModules(imported, ModuleDecl::ImportFilterKind::Public);
+  while (!imported.empty()) {
+    ModuleDecl *importedModule = std::get<1>(imported.back());
+    imported.pop_back();
+    if (!seen.insert(importedModule).second)
+      continue;
+
+    using CrossImportMap =
+        llvm::SmallDenseMap<Identifier, SmallVector<OverlayFile *, 1>>;
+    CrossImportMap &crossImports = importedModule->declaredCrossImports;
+
+    // If any of MD's cross imports declare this module as its overlay, report
+    // MD as the underlying module.
+    auto ret = std::find_if(crossImports.begin(), crossImports.end(),
+                            [&](CrossImportMap::iterator::value_type &pair) {
+      for (OverlayFile *file: std::get<1>(pair)) {
+        ArrayRef<Identifier> overlays = file->getOverlayModuleNames(
+            importedModule, SourceLoc(), std::get<0>(pair));
+        if (std::find(overlays.begin(), overlays.end(),
+                      overlayModule->getName()) != overlays.end())
+          return true;
+      }
+      return false;
+    });
+    if (ret != crossImports.end())
+      return *(declaringModuleAndBystander = {importedModule, ret->first});
 
-  populateTransitiveCrossImports(this, declaredCrossImportsTransitive);
+    if (!importedModule->hasUnderscoredNaming())
+      continue;
 
-  auto end = declaredCrossImportsTransitive.end();
-  for (auto i = declaredCrossImportsTransitive.find(overlay);
-       i != end;
-       i = declaredCrossImportsTransitive.find(i->second.second)) {
-    bystanders.push_back(i->second.first);
+    furtherImported.clear();
+    importedModule->getImportedModules(furtherImported,
+                                       ModuleDecl::ImportFilterKind::Public);
+    imported.append(furtherImported.begin(), furtherImported.end());
   }
+
+  return *(declaringModuleAndBystander = {nullptr, Identifier()});
 }
 
-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;
-  });
+bool ModuleDecl::isCrossImportOverlayOf(ModuleDecl *other) {
+  ModuleDecl *current = this;
+  while ((current = current->getDeclaringModuleAndBystander().first)) {
+    if (current == other)
+      return true;
+  }
+  return false;
+}
+
+ModuleDecl *ModuleDecl::getDeclaringModuleIfCrossImportOverlay() {
+  ModuleDecl *current = this, *declaring = nullptr;
+  while ((current = current->getDeclaringModuleAndBystander().first))
+    declaring = current;
+  return declaring;
+}
+
+bool ModuleDecl::getRequiredBystandersIfCrossImportOverlay(
+    ModuleDecl *declaring, SmallVectorImpl<Identifier> &bystanderNames) {
+  auto current = std::make_pair(this, Identifier());
+  while ((current = current.first->getDeclaringModuleAndBystander()).first) {
+    bystanderNames.push_back(current.second);
+    if (current.first == declaring)
+      return true;
+  }
+  return false;
 }
 
 namespace {
@@ -2170,32 +2210,6 @@ bool SourceFile::shouldCrossImport() const {
          getASTContext().LangOpts.EnableCrossImportOverlays;
 }
 
-ModuleDecl*
-SourceFile::getModuleShadowedBySeparatelyImportedOverlay(const ModuleDecl *overlay) {
-  if (separatelyImportedOverlaysReversed.empty() &&
-      !separatelyImportedOverlays.empty()) {
-    for (auto &entry: separatelyImportedOverlays) {
-      ModuleDecl *shadowed = entry.first;
-      for (ModuleDecl *overlay: entry.second) {
-        // If for some reason the same overlay shadows more than one module,
-        // pick the one whose name is alphabetically first.
-        ModuleDecl *old = separatelyImportedOverlaysReversed[overlay];
-        if (!old || shadowed->getNameStr() < old->getNameStr())
-          separatelyImportedOverlaysReversed[overlay] = shadowed;
-      }
-    }
-  }
-
-  ModuleDecl *underlying = const_cast<ModuleDecl *>(overlay);
-  while (underlying->getNameStr().startswith("_")) {
-    auto next = separatelyImportedOverlaysReversed.find(underlying);
-    if (next == separatelyImportedOverlaysReversed.end())
-      return nullptr;
-    underlying = std::get<1>(*next);
-  }
-  return underlying;
-};
-
 void ModuleDecl::clearLookupCache() {
   getASTContext().getImportCache().clear();