Add support for loading cross-import files

Author: beccadax

include/swift/AST/DiagnosticsCommon.def

@@ -165,6 +165,16 @@ WARNING(warn_property_wrapper_module_scope,none,
         "wrapper %0; please qualify the reference with %1",
         (DeclNameRef, Identifier))
 
+//------------------------------------------------------------------------------
+// MARK: Cross-import overlay loading diagnostics
+//------------------------------------------------------------------------------
+ERROR(cannot_load_swiftoverlay_file, none,
+      "cannot load cross-import overlay for %0 and %1: %2 (declared by '%3')",
+      (Identifier, Identifier, StringRef, StringRef))
+ERROR(cannot_list_swiftcrossimport_dir, none,
+      "cannot list cross-import overlays for %0: %1 (declared in '%2')",
+      (Identifier, StringRef, StringRef))
+
 #ifndef DIAG_NO_UNDEF
 # if defined(DIAG)
 #  undef DIAG

include/swift/AST/FileUnit.h

@@ -214,6 +214,13 @@ class FileUnit : public DeclContext {
   virtual void
   collectLinkLibraries(ModuleDecl::LinkLibraryCallback callback) const {}
 
+  /// Returns the path of the file or directory that defines the module
+  /// represented by this \c FileUnit, or empty string if there is none.
+  /// Cross-import overlay specifiers are found relative to this path.
+  virtual StringRef getModuleDefiningPath() const {
+    return "";
+  }
+
   /// True if this file contains the main class for the module.
   bool hasMainClass() const {
     return getMainClass();

include/swift/AST/Module.h

@@ -124,6 +124,8 @@ enum class ResilienceStrategy : unsigned {
   Resilient
 };
 
+class OverlayFile;
+
 /// The minimum unit of compilation.
 ///
 /// A module is made up of several file-units, which are all part of the same
@@ -203,6 +205,10 @@ class ModuleDecl : public DeclContext, public TypeDecl {
 
   SmallVector<FileUnit *, 2> Files;
 
+  friend OverlayFile;
+  llvm::SmallDenseMap<Identifier, SmallVector<OverlayFile *, 1>>
+    declaredCrossImports;
+
   std::unique_ptr<SourceLookupCache> Cache;
   SourceLookupCache &getSourceLookupCache() const;
 
@@ -252,6 +258,22 @@ class ModuleDecl : public DeclContext, public TypeDecl {
   void addFile(FileUnit &newFile);
   void removeFile(FileUnit &existingFile);
 
+  /// Add a file declaring a cross-import overlay.
+  void addCrossImportOverlayFile(StringRef file);
+
+  /// Append to \p overlayNames the names of all modules that this module
+  /// declares should be imported when \p bystanderName is imported.
+  ///
+  /// This operation is asymmetric: you will get different results if you
+  /// reverse the positions of the two modules involved in the cross-import.
+  void findDeclaredCrossImportOverlays(
+      Identifier bystanderName, SmallVectorImpl<Identifier> &overlayNames,
+      SourceLoc diagLoc);
+
+  /// Get the list of all modules this module declares a cross-import with.
+  void getDeclaredCrossImportBystanders(
+      SmallVectorImpl<Identifier> &bystanderNames);
+
   /// Convenience accessor for clients that know what kind of file they're
   /// dealing with.
   SourceFile &getMainSourceFile(SourceFileKind expectedKind) const;

include/swift/AST/ModuleLoader.h

@@ -38,6 +38,7 @@ namespace swift {
 class AbstractFunctionDecl;
 class ClangImporterOptions;
 class ClassDecl;
+class FileUnit;
 class ModuleDecl;
 class NominalTypeDecl;
 class TypeDecl;
@@ -154,6 +155,10 @@ class ModuleLoader {
 
   /// Verify all modules loaded by this loader.
   virtual void verifyAllModules() { }
+
+  /// Discover overlays declared alongside this file and add infomation about
+  /// them to it.
+  void findOverlayFiles(SourceLoc diagLoc, ModuleDecl *module, FileUnit *file);
 };
 
 } // namespace swift

include/swift/ClangImporter/ClangModule.h

@@ -118,6 +118,8 @@ class ClangModuleUnit final : public LoadedFile {
   Optional<clang::ExternalASTSource::ASTSourceDescriptor>
   getASTSourceDescriptor() const;
 
+  virtual StringRef getModuleDefiningPath() const override;
+
   static bool classof(const FileUnit *file) {
     return file->getKind() == FileUnitKind::ClangModule;
   }

include/swift/Serialization/SerializedModuleLoader.h

@@ -393,6 +393,8 @@ class SerializedASTFile final : public LoadedFile {
 
   virtual StringRef getFilename() const override;
 
+  virtual StringRef getModuleDefiningPath() const override;
+
   ClassDecl *getMainClass() const override;
 
   bool hasEntryPoint() const override;

lib/AST/Module.cpp

@@ -1454,6 +1454,189 @@ const clang::Module *ModuleDecl::findUnderlyingClangModule() const {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// Cross-Import Overlays
+//===----------------------------------------------------------------------===//
+
+namespace swift {
+/// Represents a file containing information about cross-module overlays.
+class OverlayFile {
+  /// If not empty, contains the name of the cross-import file that must be
+  /// loaded; \c overlayModuleNames is not valid and must be empty.
+  ///
+  /// If empty, \c overlayModuleNames is valid and will contain the list of
+  /// module names unless loading the file failed.
+  StringRef filePath;
+
+  /// The list of module names; empty if loading failed.
+  llvm::TinyPtrVector<Identifier> overlayModuleNames;
+
+  /// Actually loads the overlay module name list. This should mutate
+  /// \c overlayModuleNames, but not \c filePath.
+  void loadOverlayModuleNames(ModuleDecl *M, SourceLoc diagLoc,
+                              Identifier bystandingModule);
+
+public:
+  // Only allocate in ASTContexts.
+  void *operator new(size_t bytes) = delete;
+  void *operator new(size_t bytes, const ASTContext &ctx,
+                     unsigned alignment = alignof(OverlayFile)) {
+    return ctx.Allocate(bytes, alignment);
+  }
+
+  OverlayFile(StringRef filePath)
+    : filePath(filePath), overlayModuleNames() { }
+
+  /// Returns the list of additional modules that should be imported if both
+  /// the primary and secondary modules have been imported. This may load a
+  /// file; if so, it will diagnose any errors itself and arrange for the file
+  /// to not be loaded again.
+  ///
+  /// The result can be empty, either because of an error or because the file
+  /// didn't contain any overlay module names.
+  ArrayRef<Identifier> getOverlayModuleNames(ModuleDecl *M, SourceLoc diagLoc,
+                                             Identifier bystandingModule) {
+    if (!filePath.empty()) {
+      assert(overlayModuleNames.empty()
+             && "cross-import list can't be full before loading");
+      loadOverlayModuleNames(M, diagLoc, bystandingModule);
+      filePath = "";
+    }
+    return overlayModuleNames;
+  }
+};
+}
+
+void ModuleDecl::addCrossImportOverlayFile(StringRef file) {
+  auto &ctx = getASTContext();
+
+  Identifier secondaryModule = ctx.getIdentifier(llvm::sys::path::stem(file));
+  declaredCrossImports[secondaryModule]
+      .push_back(new (ctx) OverlayFile(ctx.AllocateCopy(file)));
+}
+
+void ModuleDecl::
+findDeclaredCrossImportOverlays(Identifier bystanderName,
+                                SmallVectorImpl<Identifier> &overlayNames,
+                                SourceLoc diagLoc) {
+  if (getName() == bystanderName)
+    // We don't currently support self-cross-imports.
+    return;
+
+  for (auto &crossImportFile : declaredCrossImports[bystanderName])
+    llvm::copy(crossImportFile->getOverlayModuleNames(this, diagLoc,
+                                                      bystanderName),
+               std::back_inserter(overlayNames));
+}
+
+void ModuleDecl::getDeclaredCrossImportBystanders(
+    SmallVectorImpl<Identifier> &otherModules) {
+  for (auto &pair : declaredCrossImports)
+    otherModules.push_back(std::get<0>(pair));
+}
+
+namespace {
+struct OverlayFileContents {
+  struct Module {
+    std::string name;
+  };
+
+  unsigned version;
+  std::vector<Module> modules;
+
+  static llvm::ErrorOr<OverlayFileContents>
+  load(std::unique_ptr<llvm::MemoryBuffer> input,
+       SmallVectorImpl<std::string> &errorMessages);
+};
+} // end anonymous namespace
+
+namespace llvm {
+namespace yaml {
+template <>
+struct MappingTraits<OverlayFileContents::Module> {
+  static void mapping(IO &io, OverlayFileContents::Module &module) {
+    io.mapRequired("name", module.name);
+  }
+};
+
+template <>
+struct SequenceElementTraits<OverlayFileContents::Module> {
+  static const bool flow = false;
+};
+
+template <>
+struct MappingTraits<OverlayFileContents> {
+  static void mapping(IO &io, OverlayFileContents &contents) {
+    io.mapRequired("version", contents.version);
+    io.mapRequired("modules", contents.modules);
+  }
+};
+}
+} // end namespace 'llvm'
+
+static void pushYAMLError(const llvm::SMDiagnostic &diag, void *Context) {
+  auto &errorMessages = *static_cast<SmallVectorImpl<std::string> *>(Context);
+  errorMessages.emplace_back(diag.getMessage());
+}
+
+llvm::ErrorOr<OverlayFileContents>
+OverlayFileContents::load(std::unique_ptr<llvm::MemoryBuffer> input,
+                          SmallVectorImpl<std::string> &errorMessages) {
+  llvm::yaml::Input yamlInput(input->getBuffer(), /*Ctxt=*/nullptr,
+                              pushYAMLError, &errorMessages);
+  OverlayFileContents contents;
+  yamlInput >> contents;
+
+  if (auto error = yamlInput.error())
+    return error;
+
+  if (contents.version > 1) {
+    auto message = Twine("key 'version' has invalid value: ") + Twine(contents.version);
+    errorMessages.push_back(message.str());
+    return make_error_code(std::errc::result_out_of_range);
+  }
+
+  return contents;
+}
+
+void OverlayFile::loadOverlayModuleNames(ModuleDecl *M, SourceLoc diagLoc,
+                                         Identifier bystanderName) {
+  assert(!filePath.empty());
+
+  auto &ctx = M->getASTContext();
+  llvm::vfs::FileSystem &fs = *ctx.SourceMgr.getFileSystem();
+
+  auto bufOrError = fs.getBufferForFile(filePath);
+  if (!bufOrError) {
+    M->getASTContext().Diags
+        .diagnose(diagLoc, diag::cannot_load_swiftoverlay_file,
+                  M->getName(), bystanderName, bufOrError.getError().message(),
+                  filePath);
+    return;
+  }
+
+  SmallVector<std::string, 4> errorMessages;
+  auto contentsOrErr = OverlayFileContents::load(std::move(bufOrError.get()),
+                                                 errorMessages);
+  if (!contentsOrErr) {
+    if (errorMessages.empty())
+      errorMessages.push_back(contentsOrErr.getError().message());
+
+    for (auto message : errorMessages)
+      M->getASTContext().Diags
+          .diagnose(diagLoc, diag::cannot_load_swiftoverlay_file,
+                    M->getName(), bystanderName, message, filePath);
+    return;
+  }
+
+  auto contents = std::move(*contentsOrErr);
+
+  for (const auto &module : contents.modules) {
+    auto moduleIdent = ctx.getIdentifier(module.name);
+    overlayModuleNames.push_back(moduleIdent);
+  }
+}
+
 //===----------------------------------------------------------------------===//
 // SourceFile Implementation
 //===----------------------------------------------------------------------===//

lib/AST/ModuleLoader.cpp

@@ -14,7 +14,12 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "swift/AST/ASTContext.h"
+#include "swift/AST/DiagnosticsCommon.h"
+#include "swift/AST/FileUnit.h"
 #include "swift/AST/ModuleLoader.h"
+#include "swift/Basic/FileTypes.h"
+#include "swift/Basic/Platform.h"
 #include "clang/Frontend/Utils.h"
 #include "swift/ClangImporter/ClangImporter.h"
 
@@ -54,4 +59,73 @@ DependencyTracker::getClangCollector() {
   return clangCollector;
 }
 
+static bool findOverlayFilesInDirectory(SourceLoc diagLoc, StringRef path,
+                                        ModuleDecl *module,
+                                        DependencyTracker * const tracker) {
+  using namespace llvm::sys;
+  using namespace file_types;
+
+  ASTContext &ctx = module->getASTContext();
+  auto fs = ctx.SourceMgr.getFileSystem();
+
+  std::error_code error;
+  for (auto dir = fs->dir_begin(path, error);
+       !error && dir != llvm::vfs::directory_iterator();
+       dir.increment(error)) {
+    StringRef file = dir->path();
+    if (lookupTypeForExtension(path::extension(file)) != TY_SwiftOverlayFile)
+      continue;
+
+    module->addCrossImportOverlayFile(file);
+
+    // FIXME: Better to add it only if we load it.
+    if (tracker)
+      tracker->addDependency(file, module->isSystemModule());
+  }
+
+  if (error && error != std::errc::no_such_file_or_directory) {
+    ctx.Diags.diagnose(diagLoc, diag::cannot_list_swiftcrossimport_dir,
+                       module->getName(), error.message(), path);
+  }
+  return !error;
+}
+
+void ModuleLoader::findOverlayFiles(SourceLoc diagLoc, ModuleDecl *module,
+                                    FileUnit *file) {
+  using namespace llvm::sys;
+  using namespace file_types;
+
+  auto &langOpts = module->getASTContext().LangOpts;
+
+  SmallString<64> dirPath{file->getModuleDefiningPath()};
+
+  if (dirPath.empty())
+    return;
+  path::remove_filename(dirPath);
+
+  // <parent-dir>/<module-name>.swiftcrossimport
+  path::append(dirPath, file->getExportedModuleName());
+  path::replace_extension(dirPath, getExtension(TY_SwiftCrossImportDir));
+  if (!findOverlayFilesInDirectory(diagLoc, dirPath, module, dependencyTracker))
+    // If we diagnosed an error, or we didn't find the directory at all, don't
+    // bother trying the target-specific directories.
+    return;
+
+  // <parent-dir>/<module-name>.swiftcrossimport/<target-module-triple>
+  auto moduleTriple = getTargetSpecificModuleTriple(langOpts.Target);
+  path::append(dirPath, moduleTriple.str());
+  findOverlayFilesInDirectory(diagLoc, dirPath, module, dependencyTracker);
+
+  if (!langOpts.TargetVariant)
+    return;
+
+  path::remove_filename(dirPath);
+
+  // <parent-dir>/<module-name>.swiftcrossimport/<targetvariant-module-triple>
+  auto moduleVariantTriple =
+      getTargetSpecificModuleTriple(*langOpts.TargetVariant);
+  path::append(dirPath, moduleVariantTriple.str());
+  findOverlayFilesInDirectory(diagLoc, dirPath, module, dependencyTracker);
+}
+
 } // namespace swift