Factor out the pre-registered MemoryBuffer functionality

form SerializedModuleLoader into its own ModuleLoader class. (NFC-ish)

This gives better control over the order in which the various module
load mechanisms are applied.

Author: adrian-prantl

Diff for: include/swift/ASTSectionImporter/ASTSectionImporter.h

@@ -21,18 +21,16 @@
 #include <string>
 
 namespace swift {
-  class SerializedModuleLoader;
+  class MemoryBufferSerializedModuleLoader;
 
-  /// Provided a memory buffer with an entire Mach-O __apple_ast
-  /// section, this function makes memory buffer copies of all swift
-  /// modules found in it and registers them using
-  /// registerMemoryBuffer() so they can be found by loadModule(). The
-  /// access path of all modules found in the section is appended to
-  /// the vector foundModules.
+  /// Provided a memory buffer with an entire Mach-O __swift_ast section, this
+  /// function makes memory buffer copies of all swift modules found in it and
+  /// registers them using registerMemoryBuffer() so they can be found by
+  /// loadModule(). The access path of all modules found in the section is
+  /// appended to the vector foundModules.
   /// \return true if successful.
-  bool parseASTSection(SerializedModuleLoader* SML, StringRef Data,
+  bool parseASTSection(MemoryBufferSerializedModuleLoader &Loader,
+                       StringRef Data,
                        SmallVectorImpl<std::string> &foundModules);
-
-
 }
 #endif

Diff for: include/swift/Basic/LangOptions.h

@@ -104,6 +104,10 @@ namespace swift {
     /// Enable features useful for running in the debugger.
     bool DebuggerSupport = false;
 
+    /// Enable the MemoryBufferSerializedModuleImporter.
+    /// Only used by lldb-moduleimport-test.
+    bool EnableMemoryBufferImporter = false;
+
     /// Enable the DWARFImporter. Only used by lldb-moduleimport-test.
     bool EnableDWARFImporter = false;
 

Diff for: include/swift/Frontend/Frontend.h

@@ -52,6 +52,7 @@
 namespace swift {
 
 class SerializedModuleLoader;
+class MemoryBufferSerializedModuleLoader;
 class SILModule;
 
 /// The abstract configuration of the compiler, including:
@@ -372,6 +373,7 @@ class CompilerInstance {
 
   ModuleDecl *MainModule = nullptr;
   SerializedModuleLoader *SML = nullptr;
+  MemoryBufferSerializedModuleLoader *MemoryBufferLoader = nullptr;
 
   /// Contains buffer IDs for input source code files.
   std::vector<unsigned> InputSourceCodeBufferIDs;
@@ -465,7 +467,10 @@ class CompilerInstance {
 
   ModuleDecl *getMainModule();
 
-  SerializedModuleLoader *getSerializedModuleLoader() const { return SML; }
+  MemoryBufferSerializedModuleLoader *
+  getMemoryBufferSerializedModuleLoader() const {
+    return MemoryBufferLoader;
+  }
 
   ArrayRef<unsigned> getInputBufferIDs() const {
     return InputSourceCodeBufferIDs;

Diff for: include/swift/Serialization/SerializedModuleLoader.h

@@ -39,7 +39,6 @@ class SerializedModuleLoaderBase : public ModuleLoader {
   SmallVector<std::unique_ptr<llvm::MemoryBuffer>, 2> OrphanedMemoryBuffers;
 
 protected:
-  llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> MemoryBuffers;
   ASTContext &Ctx;
   ModuleLoadingMode LoadMode;
   SerializedModuleLoaderBase(ASTContext &ctx, DependencyTracker *tracker,
@@ -170,10 +169,51 @@ class SerializedModuleLoader : public SerializedModuleLoaderBase {
 public:
   virtual ~SerializedModuleLoader();
 
-  /// Register a memory buffer that contains the serialized
-  /// module for the given access path. This API is intended to be
-  /// used by LLDB to add swiftmodules discovered in the __apple_ast
-  /// section of a Mach-O file to the search path.
+  /// Create a new importer that can load serialized Swift modules
+  /// into the given ASTContext.
+  static std::unique_ptr<SerializedModuleLoader>
+  create(ASTContext &ctx, DependencyTracker *tracker = nullptr,
+         ModuleLoadingMode loadMode = ModuleLoadingMode::PreferSerialized) {
+    return std::unique_ptr<SerializedModuleLoader>{
+      new SerializedModuleLoader(ctx, tracker, loadMode)
+    };
+  }
+};
+
+/// Imports serialized Swift modules from a MemoryBuffer into an ASTContext.
+/// This interface is primarily used by LLDB.
+class MemoryBufferSerializedModuleLoader : public SerializedModuleLoaderBase {
+  llvm::StringMap<std::unique_ptr<llvm::MemoryBuffer>> MemoryBuffers;
+
+  MemoryBufferSerializedModuleLoader(ASTContext &ctx,
+                                     DependencyTracker *tracker,
+                                     ModuleLoadingMode loadMode)
+      : SerializedModuleLoaderBase(ctx, tracker, loadMode) {}
+
+  std::error_code findModuleFilesInDirectory(
+      AccessPathElem ModuleID, StringRef DirPath, StringRef ModuleFilename,
+      StringRef ModuleDocFilename,
+      std::unique_ptr<llvm::MemoryBuffer> *ModuleBuffer,
+      std::unique_ptr<llvm::MemoryBuffer> *ModuleDocBuffer) override;
+
+  bool maybeDiagnoseTargetMismatch(SourceLoc sourceLocation,
+                                   StringRef moduleName,
+                                   StringRef archName,
+                                   StringRef directoryPath) override;
+
+public:
+  virtual ~MemoryBufferSerializedModuleLoader();
+
+  bool canImportModule(std::pair<Identifier, SourceLoc> named) override;
+  ModuleDecl *
+  loadModule(SourceLoc importLoc,
+             ArrayRef<std::pair<Identifier, SourceLoc>> path) override;
+
+  /// Register a memory buffer that contains the serialized module for the given
+  /// access path. This API is intended to be used by LLDB to add swiftmodules
+  /// discovered in the __swift_ast section of a Mach-O file (or the .swift_ast
+  /// section of an ELF file) to the search path.
+  ///
   /// FIXME: make this an actual access *path* once submodules are designed.
   void registerMemoryBuffer(StringRef AccessPath,
                             std::unique_ptr<llvm::MemoryBuffer> input) {
@@ -182,15 +222,15 @@ class SerializedModuleLoader : public SerializedModuleLoaderBase {
 
   /// Create a new importer that can load serialized Swift modules
   /// into the given ASTContext.
-  static std::unique_ptr<SerializedModuleLoader>
+  static std::unique_ptr<MemoryBufferSerializedModuleLoader>
   create(ASTContext &ctx, DependencyTracker *tracker = nullptr,
          ModuleLoadingMode loadMode = ModuleLoadingMode::PreferSerialized) {
-    return std::unique_ptr<SerializedModuleLoader>{
-      new SerializedModuleLoader(ctx, tracker, loadMode)
-    };
+    return std::unique_ptr<MemoryBufferSerializedModuleLoader>{
+        new MemoryBufferSerializedModuleLoader(ctx, tracker, loadMode)};
   }
 };
 
+
 /// A file-unit loaded from a serialized AST file.
 class SerializedASTFile final : public LoadedFile {
   friend class SerializedModuleLoaderBase;

Diff for: lib/ASTSectionImporter/ASTSectionImporter.cpp

@@ -24,7 +24,8 @@
 
 using namespace swift;
 
-bool swift::parseASTSection(SerializedModuleLoader *SML, StringRef buf,
+bool swift::parseASTSection(MemoryBufferSerializedModuleLoader &Loader,
+                            StringRef buf,
                             SmallVectorImpl<std::string> &foundModules) {
   if (!serialization::isSerializedAST(buf))
     return false;
@@ -44,7 +45,7 @@ bool swift::parseASTSection(SerializedModuleLoader *SML, StringRef buf,
           llvm::MemoryBuffer::getMemBuffer(moduleData, info.name, false));
 
         // Register the memory buffer.
-        SML->registerMemoryBuffer(info.name, std::move(bitstream));
+        Loader.registerMemoryBuffer(info.name, std::move(bitstream));
         foundModules.push_back(info.name);
       }
     } else {

Diff for: lib/Frontend/Frontend.cpp

@@ -315,6 +315,13 @@ bool CompilerInstance::setUpModuleLoaders() {
     }
   }
 
+  if (Invocation.getLangOptions().EnableMemoryBufferImporter) {
+    auto MemoryBufferLoader = MemoryBufferSerializedModuleLoader::create(
+        *Context, getDependencyTracker());
+    this->MemoryBufferLoader = MemoryBufferLoader.get();
+    Context->addModuleLoader(std::move(MemoryBufferLoader));
+  }
+
   std::unique_ptr<SerializedModuleLoader> SML =
     SerializedModuleLoader::create(*Context, getDependencyTracker(), MLM);
   this->SML = SML.get();
@@ -1067,6 +1074,7 @@ void CompilerInstance::freeASTContext() {
   Context.reset();
   MainModule = nullptr;
   SML = nullptr;
+  MemoryBufferLoader = nullptr;
   PrimaryBufferIDs.clear();
   PrimarySourceFiles.clear();
 }

Diff for: lib/Serialization/SerializedModuleLoader.cpp

@@ -35,13 +35,14 @@ namespace {
 } // end unnamed namespace
 
 // Defined out-of-line so that we can see ~ModuleFile.
-SerializedModuleLoaderBase::SerializedModuleLoaderBase(ASTContext &ctx,
-                                                       DependencyTracker *tracker,
-                                                       ModuleLoadingMode loadMode)
-  : ModuleLoader(tracker), Ctx(ctx), LoadMode(loadMode) {}
-SerializedModuleLoaderBase::~SerializedModuleLoaderBase() = default;
+SerializedModuleLoaderBase::SerializedModuleLoaderBase(
+    ASTContext &ctx, DependencyTracker *tracker, ModuleLoadingMode loadMode)
+    : ModuleLoader(tracker), Ctx(ctx), LoadMode(loadMode) {}
 
+SerializedModuleLoaderBase::~SerializedModuleLoaderBase() = default;
 SerializedModuleLoader::~SerializedModuleLoader() = default;
+MemoryBufferSerializedModuleLoader::~MemoryBufferSerializedModuleLoader() =
+    default;
 
 std::error_code SerializedModuleLoaderBase::openModuleDocFile(
   AccessPathElem ModuleID, StringRef ModuleDocPath,
@@ -607,23 +608,22 @@ void swift::serialization::diagnoseSerializedASTLoadFailure(
   }
 }
 
-bool
-SerializedModuleLoaderBase::canImportModule(std::pair<Identifier, SourceLoc> mID) {
-  // First see if we find it in the registered memory buffers.
-  if (!MemoryBuffers.empty()) {
-    auto bufIter = MemoryBuffers.find(mID.first.str());
-    if (bufIter != MemoryBuffers.end()) {
-      return true;
-    }
-  }
-
-  // Otherwise look on disk.
+bool SerializedModuleLoaderBase::canImportModule(
+    std::pair<Identifier, SourceLoc> mID) {
+  // Look on disk.
   bool isFramework = false;
   return findModule(mID, nullptr, nullptr, isFramework);
 }
 
-ModuleDecl *SerializedModuleLoaderBase::loadModule(SourceLoc importLoc,
-                                                   ModuleDecl::AccessPathTy path) {
+bool MemoryBufferSerializedModuleLoader::canImportModule(
+    std::pair<Identifier, SourceLoc> mID) {
+  // See if we find it in the registered memory buffers.
+  return MemoryBuffers.count(mID.first.str());
+}
+
+ModuleDecl *
+SerializedModuleLoaderBase::loadModule(SourceLoc importLoc,
+                                       ModuleDecl::AccessPathTy path) {
   // FIXME: Swift submodules?
   if (path.size() > 1)
     return nullptr;
@@ -633,30 +633,17 @@ ModuleDecl *SerializedModuleLoaderBase::loadModule(SourceLoc importLoc,
 
   std::unique_ptr<llvm::MemoryBuffer> moduleInputBuffer;
   std::unique_ptr<llvm::MemoryBuffer> moduleDocInputBuffer;
-  // First see if we find it in the registered memory buffers.
-  if (!MemoryBuffers.empty()) {
-    // FIXME: Right now this works only with access paths of length 1.
-    // Once submodules are designed, this needs to support suffix
-    // matching and a search path.
-    auto bufIter = MemoryBuffers.find(moduleID.first.str());
-    if (bufIter != MemoryBuffers.end()) {
-      moduleInputBuffer = std::move(bufIter->second);
-      MemoryBuffers.erase(bufIter);
-    }
-  }
 
-  // Otherwise look on disk.
-  if (!moduleInputBuffer) {
-    if (!findModule(moduleID, &moduleInputBuffer, &moduleDocInputBuffer,
-                    isFramework)) {
-      return nullptr;
-    }
-    if (dependencyTracker) {
-      // Don't record cached artifacts as dependencies.
-      StringRef DepPath = moduleInputBuffer->getBufferIdentifier();
-      if (!isCached(DepPath)) {
-        dependencyTracker->addDependency(DepPath, /*isSystem=*/false);
-      }
+  // Look on disk.
+  if (!findModule(moduleID, &moduleInputBuffer, &moduleDocInputBuffer,
+                  isFramework)) {
+    return nullptr;
+  }
+  if (dependencyTracker) {
+    // Don't record cached artifacts as dependencies.
+    StringRef DepPath = moduleInputBuffer->getBufferIdentifier();
+    if (!isCached(DepPath)) {
+      dependencyTracker->addDependency(DepPath, /*isSystem=*/false);
     }
   }
 
@@ -675,6 +662,43 @@ ModuleDecl *SerializedModuleLoaderBase::loadModule(SourceLoc importLoc,
   return M;
 }
 
+ModuleDecl *
+MemoryBufferSerializedModuleLoader::loadModule(SourceLoc importLoc,
+                                               ModuleDecl::AccessPathTy path) {
+  // FIXME: Swift submodules?
+  if (path.size() > 1)
+    return nullptr;
+
+  auto moduleID = path[0];
+
+  // See if we find it in the registered memory buffers.
+
+  // FIXME: Right now this works only with access paths of length 1.
+  // Once submodules are designed, this needs to support suffix
+  // matching and a search path.
+  auto bufIter = MemoryBuffers.find(moduleID.first.str());
+  if (bufIter == MemoryBuffers.end())
+    return nullptr;
+
+  bool isFramework = false;
+  bool treatAsPartialModule = false;
+  std::unique_ptr<llvm::MemoryBuffer> moduleInputBuffer;
+  moduleInputBuffer = std::move(bufIter->second);
+  MemoryBuffers.erase(bufIter);
+  assert(moduleInputBuffer);
+
+  auto *M = ModuleDecl::create(moduleID.first, Ctx);
+  SWIFT_DEFER { M->setHasResolvedImports(); };
+
+  if (!loadAST(*M, moduleID.second, std::move(moduleInputBuffer), {},
+               isFramework, treatAsPartialModule)) {
+    return nullptr;
+  }
+
+  Ctx.LoadedModules[moduleID.first] = M;
+  return M;
+}
+
 void SerializedModuleLoaderBase::loadExtensions(NominalTypeDecl *nominal,
                                                 unsigned previousGeneration) {
   for (auto &modulePair : LoadedModuleFiles) {
@@ -698,6 +722,24 @@ void SerializedModuleLoaderBase::loadObjCMethods(
   }
 }
 
+std::error_code MemoryBufferSerializedModuleLoader::findModuleFilesInDirectory(
+    AccessPathElem ModuleID, StringRef DirPath, StringRef ModuleFilename,
+    StringRef ModuleDocFilename,
+    std::unique_ptr<llvm::MemoryBuffer> *ModuleBuffer,
+    std::unique_ptr<llvm::MemoryBuffer> *ModuleDocBuffer) {
+  // This is a soft error instead of an llvm_unreachable because this API is
+  // primarily used by LLDB which makes it more likely that unwitting changes to
+  // the Swift compiler accidentally break the contract.
+  assert(false && "not supported");
+  return std::make_error_code(std::errc::not_supported);
+}
+
+bool MemoryBufferSerializedModuleLoader::maybeDiagnoseTargetMismatch(
+    SourceLoc sourceLocation, StringRef moduleName, StringRef archName,
+    StringRef directoryPath) {
+  return false;
+}
+
 void SerializedModuleLoaderBase::verifyAllModules() {
 #ifndef NDEBUG
   for (const LoadedModulePair &loaded : LoadedModuleFiles)

Diff for: tools/lldb-moduleimport-test/lldb-moduleimport-test.cpp

@@ -282,6 +282,7 @@ int main(int argc, char **argv) {
 
   Invocation.setModuleName("lldbtest");
   Invocation.getClangImporterOptions().ModuleCachePath = ModuleCachePath;
+  Invocation.getLangOptions().EnableMemoryBufferImporter = true;
   Invocation.getLangOptions().EnableDWARFImporter = EnableDWARFImporter;
 
   if (!ResourceDir.empty()) {
@@ -292,7 +293,7 @@ int main(int argc, char **argv) {
     return 1;
 
   for (auto &Module : Modules)
-    if (!parseASTSection(CI.getSerializedModuleLoader(),
+    if (!parseASTSection(*CI.getMemoryBufferSerializedModuleLoader(),
                          StringRef(Module.first, Module.second), modules))
       return 1;