Add mechanism for per-SourceFile overlays

When a “separately imported overlay” is added to a SourceFile, two things happen:

1. The direct import of the underlying module is removed from getImports*() by default. It is only visible if the caller passes ImportFilterKind:: ShadowedBySeparateOverlay. This means that non-module-scoped lookups will search _OverlayModule before searching its re-export UnderlyingModule, allowing it to shadow underlying declarations.

2. When you ask for lookupInModule() to look in the underlying module in that source file, it looks in the overlays instead. This means that UnderlyingModule.foo() can find declarations in _OverlayModule.

Author: beccadax

include/swift/AST/DeclContext.h

@@ -578,6 +578,14 @@ class alignas(1 << DeclContextAlignInBits) DeclContext {
                        SmallVectorImpl<ConformanceDiagnostic> *diagnostics
                          = nullptr) const;
 
+  /// Retrieves a list of separately imported overlays which are shadowing
+  /// \p declaring. If any \p overlays are returned, qualified lookups into
+  /// \p declaring should be performed into \p overlays instead; since they
+  /// are overlays, they will re-export \p declaring, but will also augment it
+  /// with additional symbols.
+  void getSeparatelyImportedOverlays(
+      ModuleDecl *declaring, SmallVectorImpl<ModuleDecl *> &overlays) const;
+
   /// Retrieve the syntactic depth of this declaration context, i.e.,
   /// the number of non-module-scoped contexts.
   ///

include/swift/AST/Module.h

@@ -466,7 +466,11 @@ class ModuleDecl : public DeclContext, public TypeDecl {
     /// Include "regular" imports with no special annotation.
     Private = 1 << 1,
     /// Include imports declared with `@_implementationOnly`.
-    ImplementationOnly = 1 << 2
+    ImplementationOnly = 1 << 2,
+    /// Include imports shadowed by a separately-imported overlay (i.e. a
+    /// cross-import overlay). Unshadowed imports are included whether or not
+    /// this flag is specified.
+    ShadowedBySeparateOverlay = 1 << 4
   };
   /// \sa getImportedModules
   using ImportFilter = OptionSet<ImportFilterKind>;

include/swift/AST/SourceFile.h

@@ -128,6 +128,17 @@ class SourceFile final : public FileUnit {
   /// The list of top-level declarations in the source file.
   std::vector<Decl *> Decls;
 
+  using SeparatelyImportedOverlayMap =
+    llvm::SmallDenseMap<ModuleDecl *, llvm::SmallPtrSet<ModuleDecl *, 1>>;
+
+  /// Keys are modules which are shadowed by one or more separately-imported
+  /// overlays; values are the list of overlays shadowing them.
+  ///
+  /// This is used by cross-import overlays to make their members appear to
+  /// be part of the underlying module. (ClangImporter overlays use a different
+  /// mechanism which is not SourceFile-dependent.)
+  SeparatelyImportedOverlayMap separatelyImportedOverlays;
+
   friend ASTContext;
   friend Impl;
 
@@ -257,6 +268,29 @@ class SourceFile final : public FileUnit {
 
   bool isImportedImplementationOnly(const ModuleDecl *module) const;
 
+  /// Register a separately-imported overlay as shadowing the module that
+  /// declares it.
+  ///
+  /// \returns true if the overlay was added; false if it already existed.
+  bool addSeparatelyImportedOverlay(ModuleDecl *overlay,
+                                    ModuleDecl *declaring) {
+    return std::get<1>(separatelyImportedOverlays[declaring].insert(overlay));
+  }
+
+  /// Retrieves a list of separately imported overlays which are shadowing
+  /// \p declaring. If any \p overlays are returned, qualified lookups into
+  /// \p declaring should be performed into \p overlays instead; since they
+  /// are overlays, they will re-export \p declaring, but will also augment it
+  /// with additional symbols.
+  void getSeparatelyImportedOverlays(
+      ModuleDecl *declaring, SmallVectorImpl<ModuleDecl *> &overlays) {
+    auto i = separatelyImportedOverlays.find(declaring);
+    if (i == separatelyImportedOverlays.end()) return;
+
+    auto &value = std::get<1>(*i);
+    overlays.append(value.begin(), value.end());
+  }
+
   void cacheVisibleDecls(SmallVectorImpl<ValueDecl *> &&globals) const;
   const SmallVectorImpl<ValueDecl *> &getCachedVisibleDecls() const;
 

lib/AST/DeclContext.cpp

@@ -287,6 +287,12 @@ DeclContext *DeclContext::getModuleScopeContext() const {
   }
 }
 
+void DeclContext::getSeparatelyImportedOverlays(
+    ModuleDecl *declaring, SmallVectorImpl<ModuleDecl *> &overlays) const {
+  if (auto SF = getParentSourceFile())
+    SF->getSeparatelyImportedOverlays(declaring, overlays);
+}
+
 /// Determine whether the given context is generic at any level.
 bool DeclContext::isGenericContext() const {
   auto dc = this;

lib/AST/Module.cpp

@@ -1162,15 +1162,18 @@ SourceFile::getImportedModules(SmallVectorImpl<ModuleDecl::ImportedModule> &modu
   assert(ASTStage >= Parsed || Kind == SourceFileKind::SIL);
   assert(filter && "no imports requested?");
   for (auto desc : Imports) {
-    ModuleDecl::ImportFilterKind requiredKind;
+    ModuleDecl::ImportFilter requiredFilter;
     if (desc.importOptions.contains(ImportFlags::Exported))
-      requiredKind = ModuleDecl::ImportFilterKind::Public;
+      requiredFilter |= ModuleDecl::ImportFilterKind::Public;
     else if (desc.importOptions.contains(ImportFlags::ImplementationOnly))
-      requiredKind = ModuleDecl::ImportFilterKind::ImplementationOnly;
+      requiredFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
     else
-      requiredKind = ModuleDecl::ImportFilterKind::Private;
+      requiredFilter |= ModuleDecl::ImportFilterKind::Private;
 
-    if (filter.contains(requiredKind))
+    if (!separatelyImportedOverlays.lookup(desc.module.second).empty())
+      requiredFilter |= ModuleDecl::ImportFilterKind::ShadowedBySeparateOverlay;
+
+    if (filter.contains(requiredFilter))
       modules.push_back(desc.module);
   }
 }

lib/AST/ModuleNameLookup.cpp

@@ -127,6 +127,20 @@ void ModuleNameLookup<LookupStrategy>::lookupInModule(
     const DeclContext *moduleScopeContext) {
   assert(moduleOrFile->isModuleScopeContext());
 
+  // Does the module scope have any separately-imported overlays shadowing
+  // the module we're looking into?
+  SmallVector<ModuleDecl *, 4> overlays;
+  moduleScopeContext->getSeparatelyImportedOverlays(
+      moduleOrFile->getParentModule(), overlays);
+  if (!overlays.empty()) {
+    // If so, look in each of those overlays.
+    for (auto overlay : overlays)
+      lookupInModule(decls, overlay, accessPath, moduleScopeContext);
+    // FIXME: This may not work gracefully if more than one of these lookups
+    // finds something.
+    return;
+  }
+
   const size_t initialCount = decls.size();
   size_t currentCount = decls.size();
 

lib/FrontendTool/FrontendTool.cpp

@@ -404,6 +404,7 @@ static bool emitLoadedModuleTraceIfNeeded(ModuleDecl *mainModule,
   ModuleDecl::ImportFilter filter = ModuleDecl::ImportFilterKind::Public;
   filter |= ModuleDecl::ImportFilterKind::Private;
   filter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
+  filter |= ModuleDecl::ImportFilterKind::ShadowedBySeparateOverlay;
   SmallVector<ModuleDecl::ImportedModule, 8> imports;
   mainModule->getImportedModules(imports, filter);
 

lib/IDE/CodeCompletion.cpp

@@ -1720,6 +1720,7 @@ class CompletionLookup final : public swift::VisibleDeclConsumer {
     ImportFilter |= ModuleDecl::ImportFilterKind::Public;
     ImportFilter |= ModuleDecl::ImportFilterKind::Private;
     ImportFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
+    // FIXME: ImportFilterKind::ShadowedBySeparateOverlay?
 
     SmallVector<ModuleDecl::ImportedModule, 16> Imported;
     SmallVector<ModuleDecl::ImportedModule, 16> FurtherImported;
@@ -5648,6 +5649,7 @@ void CodeCompletionCallbacksImpl::doneParsing() {
       ImportFilter |= ModuleDecl::ImportFilterKind::Public;
       ImportFilter |= ModuleDecl::ImportFilterKind::Private;
       ImportFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
+      // FIXME: ImportFilterKind::ShadowedBySeparateOverlay?
       SmallVector<ModuleDecl::ImportedModule, 4> Imports;
       auto *SF = CurDeclContext->getParentSourceFile();
       SF->getImportedModules(Imports, ImportFilter);

lib/Index/Index.cpp

@@ -129,6 +129,7 @@ class SourceFileOrModule {
     ImportFilter |= ModuleDecl::ImportFilterKind::Public;
     ImportFilter |= ModuleDecl::ImportFilterKind::Private;
     ImportFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
+    // FIXME: ImportFilterKind::ShadowedBySeparateOverlay?
 
     if (auto *SF = SFOrMod.dyn_cast<SourceFile *>()) {
       SF->getImportedModules(Modules, ImportFilter);
@@ -1561,6 +1562,7 @@ void IndexSwiftASTWalker::collectRecursiveModuleImports(
   ModuleDecl::ImportFilter ImportFilter;
   ImportFilter |= ModuleDecl::ImportFilterKind::Public;
   ImportFilter |= ModuleDecl::ImportFilterKind::Private;
+  // FIXME: ImportFilterKind::ShadowedBySeparateOverlay?
   SmallVector<ModuleDecl::ImportedModule, 8> Imports;
   TopMod.getImportedModules(Imports);
 

lib/Index/IndexRecord.cpp

@@ -562,6 +562,7 @@ emitDataForSwiftSerializedModule(ModuleDecl *module,
   ModuleDecl::ImportFilter importFilter;
   importFilter |= ModuleDecl::ImportFilterKind::Public;
   importFilter |= ModuleDecl::ImportFilterKind::Private;
+  // FIXME: ImportFilterKind::ShadowedBySeparateOverlay?
   SmallVector<ModuleDecl::ImportedModule, 8> imports;
   module->getImportedModules(imports, importFilter);
   StringScratchSpace moduleNameScratch;
@@ -602,6 +603,7 @@ recordSourceFileUnit(SourceFile *primarySourceFile, StringRef indexUnitToken,
   importFilter |= ModuleDecl::ImportFilterKind::Public;
   importFilter |= ModuleDecl::ImportFilterKind::Private;
   importFilter |= ModuleDecl::ImportFilterKind::ImplementationOnly;
+  // FIXME: ImportFilterKind::ShadowedBySeparateOverlay?
   SmallVector<ModuleDecl::ImportedModule, 8> imports;
   primarySourceFile->getImportedModules(imports, importFilter);
   StringScratchSpace moduleNameScratch;