vendor/clang/clang-trunk-r305145

8 files changed, 431 insertions, 169 deletions

diff --git a/lib/Frontend/ASTMerge.cpp b/lib/Frontend/ASTMerge.cpp
index 51064da270cc9..986f98ae598b9 100644
--- a/lib/Frontend/ASTMerge.cpp
+++ b/lib/Frontend/ASTMerge.cpp
@@ -21,14 +21,13 @@ ASTMergeAction::CreateASTConsumer(CompilerInstance &CI, StringRef InFile) {
   return AdaptedAction->CreateASTConsumer(CI, InFile);
 }
 
-bool ASTMergeAction::BeginSourceFileAction(CompilerInstance &CI,
-                                           StringRef Filename) {
+bool ASTMergeAction::BeginSourceFileAction(CompilerInstance &CI) {
   // FIXME: This is a hack. We need a better way to communicate the
   // AST file, compiler instance, and file name than member variables
   // of FrontendAction.
   AdaptedAction->setCurrentInput(getCurrentInput(), takeCurrentASTUnit());
   AdaptedAction->setCompilerInstance(&CI);
-  return AdaptedAction->BeginSourceFileAction(CI, Filename);
+  return AdaptedAction->BeginSourceFileAction(CI);
 }
 
 void ASTMergeAction::ExecuteAction() {
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index 01f7ca8aba9b5..8d139f2c0a6d7 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -486,6 +486,8 @@ namespace {
 class ASTInfoCollector : public ASTReaderListener {
   Preprocessor &PP;
   ASTContext &Context;
+  HeaderSearchOptions &HSOpts;
+  PreprocessorOptions &PPOpts;
   LangOptions &LangOpt;
   std::shared_ptr<TargetOptions> &TargetOpts;
   IntrusiveRefCntPtr<TargetInfo> &Target;
@@ -493,11 +495,14 @@ class ASTInfoCollector : public ASTReaderListener {
 
   bool InitializedLanguage;
 public:
-  ASTInfoCollector(Preprocessor &PP, ASTContext &Context, LangOptions &LangOpt,
+  ASTInfoCollector(Preprocessor &PP, ASTContext &Context,
+                   HeaderSearchOptions &HSOpts, PreprocessorOptions &PPOpts,
+                   LangOptions &LangOpt,
                    std::shared_ptr<TargetOptions> &TargetOpts,
                    IntrusiveRefCntPtr<TargetInfo> &Target, unsigned &Counter)
-      : PP(PP), Context(Context), LangOpt(LangOpt), TargetOpts(TargetOpts),
-        Target(Target), Counter(Counter), InitializedLanguage(false) {}
+      : PP(PP), Context(Context), HSOpts(HSOpts), PPOpts(PPOpts),
+        LangOpt(LangOpt), TargetOpts(TargetOpts), Target(Target),
+        Counter(Counter), InitializedLanguage(false) {}
 
   bool ReadLanguageOptions(const LangOptions &LangOpts, bool Complain,
                            bool AllowCompatibleDifferences) override {
@@ -511,6 +516,20 @@ public:
     return false;
   }
 
+  virtual bool ReadHeaderSearchOptions(const HeaderSearchOptions &HSOpts,
+                                       StringRef SpecificModuleCachePath,
+                                       bool Complain) override {
+    this->HSOpts = HSOpts;
+    return false;
+  }
+
+  virtual bool
+  ReadPreprocessorOptions(const PreprocessorOptions &PPOpts, bool Complain,
+                          std::string &SuggestedPredefines) override {
+    this->PPOpts = PPOpts;
+    return false;
+  }
+
   bool ReadTargetOptions(const TargetOptions &TargetOpts, bool Complain,
                          bool AllowCompatibleDifferences) override {
     // If we've already initialized the target, don't do it again.
@@ -667,6 +686,7 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
 
   ConfigureDiags(Diags, *AST, CaptureDiagnostics);
 
+  AST->LangOpts = std::make_shared<LangOptions>();
   AST->OnlyLocalDecls = OnlyLocalDecls;
   AST->CaptureDiagnostics = CaptureDiagnostics;
   AST->Diagnostics = Diags;
@@ -682,13 +702,12 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
   AST->HeaderInfo.reset(new HeaderSearch(AST->HSOpts,
                                          AST->getSourceManager(),
                                          AST->getDiagnostics(),
-                                         AST->ASTFileLangOpts,
+                                         AST->getLangOpts(),
                                          /*Target=*/nullptr));
-
-  auto PPOpts = std::make_shared<PreprocessorOptions>();
+  AST->PPOpts = std::make_shared<PreprocessorOptions>();
 
   for (const auto &RemappedFile : RemappedFiles)
-    PPOpts->addRemappedFile(RemappedFile.first, RemappedFile.second);
+    AST->PPOpts->addRemappedFile(RemappedFile.first, RemappedFile.second);
 
   // Gather Info for preprocessor construction later on.
 
@@ -696,13 +715,13 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
   unsigned Counter;
 
   AST->PP = std::make_shared<Preprocessor>(
-      std::move(PPOpts), AST->getDiagnostics(), AST->ASTFileLangOpts,
-      AST->getSourceManager(), *AST->PCMCache, HeaderInfo, *AST,
+      AST->PPOpts, AST->getDiagnostics(), *AST->LangOpts,
+      AST->getSourceManager(), *AST->PCMCache, HeaderInfo, AST->ModuleLoader,
       /*IILookup=*/nullptr,
       /*OwnsHeaderSearch=*/false);
   Preprocessor &PP = *AST->PP;
 
-  AST->Ctx = new ASTContext(AST->ASTFileLangOpts, AST->getSourceManager(),
+  AST->Ctx = new ASTContext(*AST->LangOpts, AST->getSourceManager(),
                             PP.getIdentifierTable(), PP.getSelectorTable(),
                             PP.getBuiltinInfo());
   ASTContext &Context = *AST->Ctx;
@@ -716,8 +735,8 @@ std::unique_ptr<ASTUnit> ASTUnit::LoadFromASTFile(
                               AllowPCHWithCompilerErrors);
 
   AST->Reader->setListener(llvm::make_unique<ASTInfoCollector>(
-      *AST->PP, Context, AST->ASTFileLangOpts, AST->TargetOpts, AST->Target,
-      Counter));
+      *AST->PP, Context, *AST->HSOpts, *AST->PPOpts, *AST->LangOpts,
+      AST->TargetOpts, AST->Target, Counter));
 
   // Attach the AST reader to the AST context as an external AST
   // source, so that declarations will be deserialized from the
@@ -1140,6 +1159,8 @@ bool ASTUnit::Parse(std::shared_ptr<PCHContainerOperations> PCHContainerOps,
   if (SavedMainFileBuffer)
     TranslateStoredDiagnostics(getFileManager(), getSourceManager(),
                                PreambleDiagnostics, StoredDiagnostics);
+  else
+    PreambleSrcLocCache.clear();
 
   if (!Act->Execute())
     goto error;
@@ -1963,7 +1984,7 @@ ASTUnit *ASTUnit::LoadFromCommandLine(
     unsigned PrecompilePreambleAfterNParses, TranslationUnitKind TUKind,
     bool CacheCodeCompletionResults, bool IncludeBriefCommentsInCodeCompletion,
     bool AllowPCHWithCompilerErrors, bool SkipFunctionBodies,
-    bool UserFilesAreVolatile, bool ForSerialization,
+    bool SingleFileParse, bool UserFilesAreVolatile, bool ForSerialization,
     llvm::Optional<StringRef> ModuleFormat, std::unique_ptr<ASTUnit> *ErrAST,
     IntrusiveRefCntPtr<vfs::FileSystem> VFS) {
   assert(Diags.get() && "no DiagnosticsEngine was provided");
@@ -1992,6 +2013,7 @@ ASTUnit *ASTUnit::LoadFromCommandLine(
   PPOpts.RemappedFilesKeepOriginalName = RemappedFilesKeepOriginalName;
   PPOpts.AllowPCHWithCompilerErrors = AllowPCHWithCompilerErrors;
   PPOpts.GeneratePreamble = PrecompilePreambleAfterNParses != 0;
+  PPOpts.SingleFileParseMode = SingleFileParse;
   
   // Override the resources path.
   CI->getHeaderSearchOpts().ResourceDir = ResourceFilesPath;
@@ -2582,11 +2604,9 @@ void ASTUnit::TranslateStoredDiagnostics(
   // remap all the locations to the new view. This includes the diag location,
   // any associated source ranges, and the source ranges of associated fix-its.
   // FIXME: There should be a cleaner way to do this.
-
   SmallVector<StoredDiagnostic, 4> Result;
   Result.reserve(Diags.size());
-  const FileEntry *PreviousFE = nullptr;
-  FileID FID;
+
   for (const StandaloneDiagnostic &SD : Diags) {
     // Rebuild the StoredDiagnostic.
     if (SD.Filename.empty())
@@ -2594,11 +2614,16 @@ void ASTUnit::TranslateStoredDiagnostics(
     const FileEntry *FE = FileMgr.getFile(SD.Filename);
     if (!FE)
       continue;
-    if (FE != PreviousFE) {
-      FID = SrcMgr.translateFile(FE);
-      PreviousFE = FE;
+    SourceLocation FileLoc;
+    auto ItFileID = PreambleSrcLocCache.find(SD.Filename);
+    if (ItFileID == PreambleSrcLocCache.end()) {
+      FileID FID = SrcMgr.translateFile(FE);
+      FileLoc = SrcMgr.getLocForStartOfFile(FID);
+      PreambleSrcLocCache[SD.Filename] = FileLoc;
+    } else {
+      FileLoc = ItFileID->getValue();
     }
-    SourceLocation FileLoc = SrcMgr.getLocForStartOfFile(FID);
+
     if (FileLoc.isInvalid())
       continue;
     SourceLocation L = FileLoc.getLocWithOffset(SD.LocOffset);
@@ -2879,7 +2904,32 @@ const FileEntry *ASTUnit::getPCHFile() {
 }
 
 bool ASTUnit::isModuleFile() {
-  return isMainFileAST() && ASTFileLangOpts.isCompilingModule();
+  return isMainFileAST() && getLangOpts().isCompilingModule();
+}
+
+InputKind ASTUnit::getInputKind() const {
+  auto &LangOpts = getLangOpts();
+
+  InputKind::Language Lang;
+  if (LangOpts.OpenCL)
+    Lang = InputKind::OpenCL;
+  else if (LangOpts.CUDA)
+    Lang = InputKind::CUDA;
+  else if (LangOpts.RenderScript)
+    Lang = InputKind::RenderScript;
+  else if (LangOpts.CPlusPlus)
+    Lang = LangOpts.ObjC1 ? InputKind::ObjCXX : InputKind::CXX;
+  else
+    Lang = LangOpts.ObjC1 ? InputKind::ObjC : InputKind::C;
+
+  InputKind::Format Fmt = InputKind::Source;
+  if (LangOpts.getCompilingModule() == LangOptions::CMK_ModuleMap)
+    Fmt = InputKind::ModuleMap;
+
+  // We don't know if input was preprocessed. Assume not.
+  bool PP = false;
+
+  return InputKind(Lang, Fmt, PP);
 }
 
 void ASTUnit::PreambleData::countLines() const {
diff --git a/lib/Frontend/CompilerInstance.cpp b/lib/Frontend/CompilerInstance.cpp
index e92672a785daf..72a8c38180931 100644
--- a/lib/Frontend/CompilerInstance.cpp
+++ b/lib/Frontend/CompilerInstance.cpp
@@ -667,6 +667,11 @@ void CompilerInstance::clearOutputFiles(bool EraseFiles) {
       llvm::sys::fs::remove(OF.Filename);
   }
   OutputFiles.clear();
+  if (DeleteBuiltModules) {
+    for (auto &Module : BuiltModules)
+      llvm::sys::fs::remove(Module.second);
+    BuiltModules.clear();
+  }
   NonSeekStream.reset();
 }
 
@@ -1029,13 +1034,14 @@ static InputKind::Language getLanguageFromOptions(const LangOptions &LangOpts) {
 /// \brief Compile a module file for the given module, using the options 
 /// provided by the importing compiler instance. Returns true if the module
 /// was built without errors.
-static bool compileModuleImpl(CompilerInstance &ImportingInstance,
-                              SourceLocation ImportLoc,
-                              Module *Module,
-                              StringRef ModuleFileName) {
-  ModuleMap &ModMap 
-    = ImportingInstance.getPreprocessor().getHeaderSearchInfo().getModuleMap();
-    
+static bool
+compileModuleImpl(CompilerInstance &ImportingInstance, SourceLocation ImportLoc,
+                  StringRef ModuleName, FrontendInputFile Input,
+                  StringRef OriginalModuleMapFile, StringRef ModuleFileName,
+                  llvm::function_ref<void(CompilerInstance &)> PreBuildStep =
+                      [](CompilerInstance &) {},
+                  llvm::function_ref<void(CompilerInstance &)> PostBuildStep =
+                      [](CompilerInstance &) {}) {
   // Construct a compiler invocation for creating this module.
   auto Invocation =
       std::make_shared<CompilerInvocation>(ImportingInstance.getInvocation());
@@ -1060,7 +1066,7 @@ static bool compileModuleImpl(CompilerInstance &ImportingInstance,
       PPOpts.Macros.end());
 
   // Note the name of the module we're building.
-  Invocation->getLangOpts()->CurrentModule = Module->getTopLevelModuleName();
+  Invocation->getLangOpts()->CurrentModule = ModuleName;
 
   // Make sure that the failed-module structure has been allocated in
   // the importing instance, and propagate the pointer to the newly-created
@@ -1080,13 +1086,10 @@ static bool compileModuleImpl(CompilerInstance &ImportingInstance,
   FrontendOpts.DisableFree = false;
   FrontendOpts.GenerateGlobalModuleIndex = false;
   FrontendOpts.BuildingImplicitModule = true;
-  FrontendOpts.OriginalModuleMap =
-      ModMap.getModuleMapFileForUniquing(Module)->getName();
+  FrontendOpts.OriginalModuleMap = OriginalModuleMapFile;
   // Force implicitly-built modules to hash the content of the module file.
   HSOpts.ModulesHashContent = true;
-  FrontendOpts.Inputs.clear();
-  InputKind IK(getLanguageFromOptions(*Invocation->getLangOpts()),
-               InputKind::ModuleMap);
+  FrontendOpts.Inputs = {Input};
 
   // Don't free the remapped file buffers; they are owned by our caller.
   PPOpts.RetainRemappedFileBuffers = true;
@@ -1117,7 +1120,7 @@ static bool compileModuleImpl(CompilerInstance &ImportingInstance,
   SourceManager &SourceMgr = Instance.getSourceManager();
   SourceMgr.setModuleBuildStack(
     ImportingInstance.getSourceManager().getModuleBuildStack());
-  SourceMgr.pushModuleBuildStack(Module->getTopLevelModuleName(),
+  SourceMgr.pushModuleBuildStack(ModuleName,
     FullSourceLoc(ImportLoc, ImportingInstance.getSourceManager()));
 
   // If we're collecting module dependencies, we need to share a collector
@@ -1126,32 +1129,11 @@ static bool compileModuleImpl(CompilerInstance &ImportingInstance,
   Instance.setModuleDepCollector(ImportingInstance.getModuleDepCollector());
   Inv.getDependencyOutputOpts() = DependencyOutputOptions();
 
-  // Get or create the module map that we'll use to build this module.
-  std::string InferredModuleMapContent;
-  if (const FileEntry *ModuleMapFile =
-          ModMap.getContainingModuleMapFile(Module)) {
-    // Use the module map where this module resides.
-    FrontendOpts.Inputs.emplace_back(ModuleMapFile->getName(), IK,
-                                     +Module->IsSystem);
-  } else {
-    SmallString<128> FakeModuleMapFile(Module->Directory->getName());
-    llvm::sys::path::append(FakeModuleMapFile, "__inferred_module.map");
-    FrontendOpts.Inputs.emplace_back(FakeModuleMapFile, IK, +Module->IsSystem);
-
-    llvm::raw_string_ostream OS(InferredModuleMapContent);
-    Module->print(OS);
-    OS.flush();
-
-    std::unique_ptr<llvm::MemoryBuffer> ModuleMapBuffer =
-        llvm::MemoryBuffer::getMemBuffer(InferredModuleMapContent);
-    ModuleMapFile = Instance.getFileManager().getVirtualFile(
-        FakeModuleMapFile, InferredModuleMapContent.size(), 0);
-    SourceMgr.overrideFileContents(ModuleMapFile, std::move(ModuleMapBuffer));
-  }
-
   ImportingInstance.getDiagnostics().Report(ImportLoc,
                                             diag::remark_module_build)
-    << Module->Name << ModuleFileName;
+    << ModuleName << ModuleFileName;
+
+  PreBuildStep(Instance);
 
   // Execute the action to actually build the module in-place. Use a separate
   // thread so that we get a stack large enough.
@@ -1164,9 +1146,11 @@ static bool compileModuleImpl(CompilerInstance &ImportingInstance,
       },
       ThreadStackSize);
 
+  PostBuildStep(Instance);
+
   ImportingInstance.getDiagnostics().Report(ImportLoc,
                                             diag::remark_module_build_done)
-    << Module->Name;
+    << ModuleName;
 
   // Delete the temporary module map file.
   // FIXME: Even though we're executing under crash protection, it would still
@@ -1174,13 +1158,66 @@ static bool compileModuleImpl(CompilerInstance &ImportingInstance,
   // doesn't make sense for all clients, so clean this up manually.
   Instance.clearOutputFiles(/*EraseFiles=*/true);
 
+  return !Instance.getDiagnostics().hasErrorOccurred();
+}
+
+/// \brief Compile a module file for the given module, using the options 
+/// provided by the importing compiler instance. Returns true if the module
+/// was built without errors.
+static bool compileModuleImpl(CompilerInstance &ImportingInstance,
+                              SourceLocation ImportLoc,
+                              Module *Module,
+                              StringRef ModuleFileName) {
+  InputKind IK(getLanguageFromOptions(ImportingInstance.getLangOpts()),
+               InputKind::ModuleMap);
+
+  // Get or create the module map that we'll use to build this module.
+  ModuleMap &ModMap 
+    = ImportingInstance.getPreprocessor().getHeaderSearchInfo().getModuleMap();
+  bool Result;
+  if (const FileEntry *ModuleMapFile =
+          ModMap.getContainingModuleMapFile(Module)) {
+    // Use the module map where this module resides.
+    Result = compileModuleImpl(
+        ImportingInstance, ImportLoc, Module->getTopLevelModuleName(),
+        FrontendInputFile(ModuleMapFile->getName(), IK, +Module->IsSystem),
+        ModMap.getModuleMapFileForUniquing(Module)->getName(),
+        ModuleFileName);
+  } else {
+    // FIXME: We only need to fake up an input file here as a way of
+    // transporting the module's directory to the module map parser. We should
+    // be able to do that more directly, and parse from a memory buffer without
+    // inventing this file.
+    SmallString<128> FakeModuleMapFile(Module->Directory->getName());
+    llvm::sys::path::append(FakeModuleMapFile, "__inferred_module.map");
+
+    std::string InferredModuleMapContent;
+    llvm::raw_string_ostream OS(InferredModuleMapContent);
+    Module->print(OS);
+    OS.flush();
+
+    Result = compileModuleImpl(
+        ImportingInstance, ImportLoc, Module->getTopLevelModuleName(),
+        FrontendInputFile(FakeModuleMapFile, IK, +Module->IsSystem),
+        ModMap.getModuleMapFileForUniquing(Module)->getName(),
+        ModuleFileName,
+        [&](CompilerInstance &Instance) {
+      std::unique_ptr<llvm::MemoryBuffer> ModuleMapBuffer =
+          llvm::MemoryBuffer::getMemBuffer(InferredModuleMapContent);
+      ModuleMapFile = Instance.getFileManager().getVirtualFile(
+          FakeModuleMapFile, InferredModuleMapContent.size(), 0);
+      Instance.getSourceManager().overrideFileContents(
+          ModuleMapFile, std::move(ModuleMapBuffer));
+    });
+  }
+
   // We've rebuilt a module. If we're allowed to generate or update the global
   // module index, record that fact in the importing compiler instance.
   if (ImportingInstance.getFrontendOpts().GenerateGlobalModuleIndex) {
     ImportingInstance.setBuildGlobalModuleIndex(true);
   }
 
-  return !Instance.getDiagnostics().hasErrorOccurred();
+  return Result;
 }
 
 static bool compileAndLoadModule(CompilerInstance &ImportingInstance,
@@ -1586,24 +1623,36 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
         PP->getHeaderSearchInfo().getHeaderSearchOpts();
 
     std::string ModuleFileName;
-    bool LoadFromPrebuiltModulePath = false;
-    // We try to load the module from the prebuilt module paths. If not
-    // successful, we then try to find it in the module cache.
-    if (!HSOpts.PrebuiltModulePaths.empty()) {
-      // Load the module from the prebuilt module path.
+    enum ModuleSource {
+      ModuleNotFound, ModuleCache, PrebuiltModulePath, ModuleBuildPragma
+    } Source = ModuleNotFound;
+
+    // Check to see if the module has been built as part of this compilation
+    // via a module build pragma.
+    auto BuiltModuleIt = BuiltModules.find(ModuleName);
+    if (BuiltModuleIt != BuiltModules.end()) {
+      ModuleFileName = BuiltModuleIt->second;
+      Source = ModuleBuildPragma;
+    }
+
+    // Try to load the module from the prebuilt module path.
+    if (Source == ModuleNotFound && !HSOpts.PrebuiltModulePaths.empty()) {
       ModuleFileName = PP->getHeaderSearchInfo().getModuleFileName(
           ModuleName, "", /*UsePrebuiltPath*/ true);
       if (!ModuleFileName.empty())
-        LoadFromPrebuiltModulePath = true;
+        Source = PrebuiltModulePath;
     }
-    if (!LoadFromPrebuiltModulePath && Module) {
-      // Load the module from the module cache.
+
+    // Try to load the module from the module cache.
+    if (Source == ModuleNotFound && Module) {
       ModuleFileName = PP->getHeaderSearchInfo().getModuleFileName(Module);
-    } else if (!LoadFromPrebuiltModulePath) {
+      Source = ModuleCache;
+    }
+
+    if (Source == ModuleNotFound) {
       // We can't find a module, error out here.
       getDiagnostics().Report(ModuleNameLoc, diag::err_module_not_found)
-      << ModuleName
-      << SourceRange(ImportLoc, ModuleNameLoc);
+          << ModuleName << SourceRange(ImportLoc, ModuleNameLoc);
       ModuleBuildFailed = true;
       return ModuleLoadResult();
     }
@@ -1631,20 +1680,20 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
                  *FrontendTimerGroup);
     llvm::TimeRegion TimeLoading(FrontendTimerGroup ? &Timer : nullptr);
 
-    // Try to load the module file. If we are trying to load from the prebuilt
-    // module path, we don't have the module map files and don't know how to
-    // rebuild modules.
-    unsigned ARRFlags = LoadFromPrebuiltModulePath ?
-                        ASTReader::ARR_ConfigurationMismatch :
-                        ASTReader::ARR_OutOfDate | ASTReader::ARR_Missing;
+    // Try to load the module file. If we are not trying to load from the
+    // module cache, we don't know how to rebuild modules.
+    unsigned ARRFlags = Source == ModuleCache ?
+                        ASTReader::ARR_OutOfDate | ASTReader::ARR_Missing :
+                        ASTReader::ARR_ConfigurationMismatch;
     switch (ModuleManager->ReadAST(ModuleFileName,
-                                   LoadFromPrebuiltModulePath ?
-                                   serialization::MK_PrebuiltModule :
-                                   serialization::MK_ImplicitModule,
-                                   ImportLoc,
-                                   ARRFlags)) {
+                                   Source == PrebuiltModulePath
+                                       ? serialization::MK_PrebuiltModule
+                                       : Source == ModuleBuildPragma
+                                             ? serialization::MK_ExplicitModule
+                                             : serialization::MK_ImplicitModule,
+                                   ImportLoc, ARRFlags)) {
     case ASTReader::Success: {
-      if (LoadFromPrebuiltModulePath && !Module) {
+      if (Source != ModuleCache && !Module) {
         Module = PP->getHeaderSearchInfo().lookupModule(ModuleName);
         if (!Module || !Module->getASTFile() ||
             FileMgr->getFile(ModuleFileName) != Module->getASTFile()) {
@@ -1662,10 +1711,10 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
 
     case ASTReader::OutOfDate:
     case ASTReader::Missing: {
-      if (LoadFromPrebuiltModulePath) {
-        // We can't rebuild the module without a module map. Since ReadAST
-        // already produces diagnostics for these two cases, we simply
-        // error out here.
+      if (Source != ModuleCache) {
+        // We don't know the desired configuration for this module and don't
+        // necessarily even have a module map. Since ReadAST already produces
+        // diagnostics for these two cases, we simply error out here.
         ModuleBuildFailed = true;
         KnownModules[Path[0].first] = nullptr;
         return ModuleLoadResult();
@@ -1722,7 +1771,9 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
     }
 
     case ASTReader::ConfigurationMismatch:
-      if (LoadFromPrebuiltModulePath)
+      if (Source == PrebuiltModulePath)
+        // FIXME: We shouldn't be setting HadFatalFailure below if we only
+        // produce a warning here!
         getDiagnostics().Report(SourceLocation(),
                                 diag::warn_module_config_mismatch)
             << ModuleFileName;
@@ -1751,7 +1802,7 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
   // If we never found the module, fail.
   if (!Module)
     return ModuleLoadResult();
-  
+
   // Verify that the rest of the module path actually corresponds to
   // a submodule.
   if (Path.size() > 1) {
@@ -1824,20 +1875,10 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
     }
 
     // Check whether this module is available.
-    clang::Module::Requirement Requirement;
-    clang::Module::UnresolvedHeaderDirective MissingHeader;
-    if (!Module->isAvailable(getLangOpts(), getTarget(), Requirement,
-                             MissingHeader)) {
-      if (MissingHeader.FileNameLoc.isValid()) {
-        getDiagnostics().Report(MissingHeader.FileNameLoc,
-                                diag::err_module_header_missing)
-          << MissingHeader.IsUmbrella << MissingHeader.FileName;
-      } else {
-        getDiagnostics().Report(ImportLoc, diag::err_module_unavailable)
-          << Module->getFullModuleName()
-          << Requirement.second << Requirement.first
-          << SourceRange(Path.front().second, Path.back().second);
-      }
+    if (Preprocessor::checkModuleIsAvailable(getLangOpts(), getTarget(),
+                                             getDiagnostics(), Module)) {
+      getDiagnostics().Report(ImportLoc, diag::note_module_import_here)
+        << SourceRange(Path.front().second, Path.back().second);
       LastModuleImportLoc = ImportLoc;
       LastModuleImportResult = ModuleLoadResult();
       return ModuleLoadResult();
@@ -1858,6 +1899,53 @@ CompilerInstance::loadModule(SourceLocation ImportLoc,
   return LastModuleImportResult;
 }
 
+void CompilerInstance::loadModuleFromSource(SourceLocation ImportLoc,
+                                            StringRef ModuleName,
+                                            StringRef Source) {
+  // FIXME: Using a randomized filename here means that our intermediate .pcm
+  // output is nondeterministic (as .pcm files refer to each other by name).
+  // Can this affect the output in any way?
+  SmallString<128> ModuleFileName;
+  if (std::error_code EC = llvm::sys::fs::createTemporaryFile(
+          ModuleName, "pcm", ModuleFileName)) {
+    getDiagnostics().Report(ImportLoc, diag::err_fe_unable_to_open_output)
+        << ModuleFileName << EC.message();
+    return;
+  }
+  std::string ModuleMapFileName = (ModuleName + ".map").str();
+
+  FrontendInputFile Input(
+      ModuleMapFileName,
+      InputKind(getLanguageFromOptions(*Invocation->getLangOpts()),
+                InputKind::ModuleMap, /*Preprocessed*/true));
+
+  std::string NullTerminatedSource(Source.str());
+
+  auto PreBuildStep = [&](CompilerInstance &Other) {
+    // Create a virtual file containing our desired source.
+    // FIXME: We shouldn't need to do this.
+    const FileEntry *ModuleMapFile = Other.getFileManager().getVirtualFile(
+        ModuleMapFileName, NullTerminatedSource.size(), 0);
+    Other.getSourceManager().overrideFileContents(
+        ModuleMapFile,
+        llvm::MemoryBuffer::getMemBuffer(NullTerminatedSource.c_str()));
+
+    Other.BuiltModules = std::move(BuiltModules);
+    Other.DeleteBuiltModules = false;
+  };
+
+  auto PostBuildStep = [this](CompilerInstance &Other) {
+    BuiltModules = std::move(Other.BuiltModules);
+  };
+
+  // Build the module, inheriting any modules that we've built locally.
+  if (compileModuleImpl(*this, ImportLoc, ModuleName, Input, StringRef(),
+                        ModuleFileName, PreBuildStep, PostBuildStep)) {
+    BuiltModules[ModuleName] = ModuleFileName.str();
+    llvm::sys::RemoveFileOnSignal(ModuleFileName);
+  }
+}
+
 void CompilerInstance::makeModuleVisible(Module *Mod,
                                          Module::NameVisibilityKind Visibility,
                                          SourceLocation ImportLoc) {
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
index adb15f1730bf4..bb635b7ad714e 100644
--- a/lib/Frontend/CompilerInvocation.cpp
+++ b/lib/Frontend/CompilerInvocation.cpp
@@ -739,7 +739,9 @@ static bool ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args, InputKind IK,
   Opts.InstrumentForProfiling = Args.hasArg(OPT_pg);
   Opts.CallFEntry = Args.hasArg(OPT_mfentry);
   Opts.EmitOpenCLArgMetadata = Args.hasArg(OPT_cl_kernel_arg_info);
-  Opts.CompressDebugSections = Args.hasArg(OPT_compress_debug_sections);
+  // TODO: map this from -gz in the driver and give it a named value
+  if (Args.hasArg(OPT_compress_debug_sections))
+    Opts.setCompressDebugSections(llvm::DebugCompressionType::GNU);
   Opts.RelaxELFRelocations = Args.hasArg(OPT_mrelax_relocations);
   Opts.DebugCompilationDir = Args.getLastArgValue(OPT_fdebug_compilation_dir);
   for (auto A : Args.filtered(OPT_mlink_bitcode_file, OPT_mlink_cuda_bitcode)) {
@@ -2499,6 +2501,7 @@ static void ParsePreprocessorOutputArgs(PreprocessorOutputOptions &Opts,
   Opts.ShowMacros = Args.hasArg(OPT_dM) || Args.hasArg(OPT_dD);
   Opts.ShowIncludeDirectives = Args.hasArg(OPT_dI);
   Opts.RewriteIncludes = Args.hasArg(OPT_frewrite_includes);
+  Opts.RewriteImports = Args.hasArg(OPT_frewrite_imports);
   Opts.UseLineDirectives = Args.hasArg(OPT_fuse_line_directives);
 }
 
diff --git a/lib/Frontend/FrontendAction.cpp b/lib/Frontend/FrontendAction.cpp
index e2fbe965349f7..7c0b854648bdd 100644
--- a/lib/Frontend/FrontendAction.cpp
+++ b/lib/Frontend/FrontendAction.cpp
@@ -387,8 +387,7 @@ static std::error_code collectModuleHeaderIncludes(
   return std::error_code();
 }
 
-static bool loadModuleMapForModuleBuild(CompilerInstance &CI,
-                                        StringRef Filename, bool IsSystem,
+static bool loadModuleMapForModuleBuild(CompilerInstance &CI, bool IsSystem,
                                         bool IsPreprocessed,
                                         std::string &PresumedModuleMapFile,
                                         unsigned &Offset) {
@@ -444,21 +443,9 @@ static Module *prepareToBuildModule(CompilerInstance &CI,
   }
 
   // Check whether we can build this module at all.
-  clang::Module::Requirement Requirement;
-  clang::Module::UnresolvedHeaderDirective MissingHeader;
-  if (!M->isAvailable(CI.getLangOpts(), CI.getTarget(), Requirement,
-                      MissingHeader)) {
-    if (MissingHeader.FileNameLoc.isValid()) {
-      CI.getDiagnostics().Report(MissingHeader.FileNameLoc,
-                                 diag::err_module_header_missing)
-        << MissingHeader.IsUmbrella << MissingHeader.FileName;
-    } else {
-      CI.getDiagnostics().Report(diag::err_module_unavailable)
-        << M->getFullModuleName() << Requirement.second << Requirement.first;
-    }
-
+  if (Preprocessor::checkModuleIsAvailable(CI.getLangOpts(), CI.getTarget(),
+                                           CI.getDiagnostics(), M))
     return nullptr;
-  }
 
   // Inform the preprocessor that includes from within the input buffer should
   // be resolved relative to the build directory of the module map file.
@@ -523,7 +510,8 @@ getInputBufferForModule(CompilerInstance &CI, Module *M) {
 }
 
 bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
-                                     const FrontendInputFile &Input) {
+                                     const FrontendInputFile &RealInput) {
+  FrontendInputFile Input(RealInput);
   assert(!Instance && "Already processing a source file!");
   assert(!Input.isEmpty() && "Unexpected empty filename!");
   setCurrentInput(Input);
@@ -531,15 +519,72 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
 
   StringRef InputFile = Input.getFile();
   bool HasBegunSourceFile = false;
+  bool ReplayASTFile = Input.getKind().getFormat() == InputKind::Precompiled &&
+                       usesPreprocessorOnly();
   if (!BeginInvocation(CI))
     goto failure;
 
+  // If we're replaying the build of an AST file, import it and set up
+  // the initial state from its build.
+  if (ReplayASTFile) {
+    IntrusiveRefCntPtr<DiagnosticsEngine> Diags(&CI.getDiagnostics());
+
+    // The AST unit populates its own diagnostics engine rather than ours.
+    IntrusiveRefCntPtr<DiagnosticsEngine> ASTDiags(
+        new DiagnosticsEngine(Diags->getDiagnosticIDs(),
+                              &Diags->getDiagnosticOptions()));
+    ASTDiags->setClient(Diags->getClient(), /*OwnsClient*/false);
+
+    std::unique_ptr<ASTUnit> AST = ASTUnit::LoadFromASTFile(
+        InputFile, CI.getPCHContainerReader(), ASTDiags, CI.getFileSystemOpts(),
+        CI.getCodeGenOpts().DebugTypeExtRefs);
+    if (!AST)
+      goto failure;
+
+    // Options relating to how we treat the input (but not what we do with it)
+    // are inherited from the AST unit.
+    CI.getHeaderSearchOpts() = AST->getHeaderSearchOpts();
+    CI.getPreprocessorOpts() = AST->getPreprocessorOpts();
+    CI.getLangOpts() = AST->getLangOpts();
+
+    // Preload all the module files loaded transitively by the AST unit.
+    if (auto ASTReader = AST->getASTReader()) {
+      auto &MM = ASTReader->getModuleManager();
+      for (ModuleFile &MF : MM)
+        if (&MF != &MM.getPrimaryModule())
+          CI.getFrontendOpts().ModuleFiles.push_back(MF.FileName);
+    }
+    // FIXME: Preload module maps loaded by the AST unit.
+
+    // Set the shared objects, these are reset when we finish processing the
+    // file, otherwise the CompilerInstance will happily destroy them.
+    CI.setFileManager(&AST->getFileManager());
+    CI.createSourceManager(CI.getFileManager());
+    CI.getSourceManager().initializeForReplay(AST->getSourceManager());
+    CI.createPreprocessor(getTranslationUnitKind());
+
+    // Set up the input file for replay purposes.
+    auto Kind = AST->getInputKind();
+    if (Kind.getFormat() == InputKind::ModuleMap) {
+      Module *ASTModule =
+          AST->getPreprocessor().getHeaderSearchInfo().lookupModule(
+              AST->getLangOpts().CurrentModule, /*AllowSearch*/ false);
+      Input = FrontendInputFile(ASTModule->PresumedModuleMapFile, Kind);
+    } else {
+      auto &SM = CI.getSourceManager();
+      FileID ID = SM.getMainFileID();
+      if (auto *File = SM.getFileEntryForID(ID))
+        Input = FrontendInputFile(File->getName(), Kind);
+      else
+        Input = FrontendInputFile(SM.getBuffer(ID), Kind);
+    }
+    setCurrentInput(Input, std::move(AST));
+  }
+
   // AST files follow a very different path, since they share objects via the
   // AST unit.
   if (Input.getKind().getFormat() == InputKind::Precompiled) {
-    // FIXME: We should not be asserting on bad command-line arguments.
-    assert(!usesPreprocessorOnly() &&
-           "Attempt to pass AST file to preprocessor only action!");
+    assert(!usesPreprocessorOnly() && "this case was handled above");
     assert(hasASTFileSupport() &&
            "This action does not have AST file support!");
 
@@ -569,7 +614,7 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
     setCurrentInput(Input, std::move(AST));
 
     // Initialize the action.
-    if (!BeginSourceFileAction(CI, InputFile))
+    if (!BeginSourceFileAction(CI))
       goto failure;
 
     // Create the AST consumer.
@@ -616,7 +661,7 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
     HasBegunSourceFile = true;
 
     // Initialize the action.
-    if (!BeginSourceFileAction(CI, InputFile))
+    if (!BeginSourceFileAction(CI))
       goto failure;
 
     // Initialize the main file entry.
@@ -680,7 +725,7 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
 
     std::string PresumedModuleMapFile;
     unsigned OffsetToContents;
-    if (loadModuleMapForModuleBuild(CI, Input.getFile(), Input.isSystem(),
+    if (loadModuleMapForModuleBuild(CI, Input.isSystem(),
                                     Input.isPreprocessed(),
                                     PresumedModuleMapFile, OffsetToContents))
       goto failure;
@@ -709,7 +754,7 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
   }
 
   // Initialize the action.
-  if (!BeginSourceFileAction(CI, InputFile))
+  if (!BeginSourceFileAction(CI))
     goto failure;
 
   // Create the AST context and consumer unless this is a preprocessor only
@@ -829,14 +874,7 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
 
   // If we failed, reset state since the client will not end up calling the
   // matching EndSourceFile().
-  failure:
-  if (isCurrentFileAST()) {
-    CI.setASTContext(nullptr);
-    CI.setPreprocessor(nullptr);
-    CI.setSourceManager(nullptr);
-    CI.setFileManager(nullptr);
-  }
-
+failure:
   if (HasBegunSourceFile)
     CI.getDiagnosticClient().EndSourceFile();
   CI.clearOutputFiles(/*EraseFiles=*/true);
@@ -914,6 +952,7 @@ void FrontendAction::EndSourceFile() {
       CI.resetAndLeakPreprocessor();
       CI.resetAndLeakSourceManager();
       CI.resetAndLeakFileManager();
+      BuryPointer(CurrentASTUnit.release());
     } else {
       CI.setPreprocessor(nullptr);
       CI.setSourceManager(nullptr);
@@ -973,11 +1012,10 @@ WrapperFrontendAction::CreateASTConsumer(CompilerInstance &CI,
 bool WrapperFrontendAction::BeginInvocation(CompilerInstance &CI) {
   return WrappedAction->BeginInvocation(CI);
 }
-bool WrapperFrontendAction::BeginSourceFileAction(CompilerInstance &CI,
-                                                  StringRef Filename) {
+bool WrapperFrontendAction::BeginSourceFileAction(CompilerInstance &CI) {
   WrappedAction->setCurrentInput(getCurrentInput());
   WrappedAction->setCompilerInstance(&CI);
-  auto Ret = WrappedAction->BeginSourceFileAction(CI, Filename);
+  auto Ret = WrappedAction->BeginSourceFileAction(CI);
   // BeginSourceFileAction may change CurrentInput, e.g. during module builds.
   setCurrentInput(WrappedAction->getCurrentInput());
   return Ret;
diff --git a/lib/Frontend/FrontendActions.cpp b/lib/Frontend/FrontendActions.cpp
index 89ac385ca45d5..c5567a09636b5 100644
--- a/lib/Frontend/FrontendActions.cpp
+++ b/lib/Frontend/FrontendActions.cpp
@@ -134,8 +134,7 @@ bool GeneratePCHAction::shouldEraseOutputFiles() {
   return ASTFrontendAction::shouldEraseOutputFiles();
 }
 
-bool GeneratePCHAction::BeginSourceFileAction(CompilerInstance &CI,
-                                              StringRef Filename) {
+bool GeneratePCHAction::BeginSourceFileAction(CompilerInstance &CI) {
   CI.getLangOpts().CompilingPCH = true;
   return true;
 }
@@ -164,11 +163,6 @@ GenerateModuleAction::CreateASTConsumer(CompilerInstance &CI,
   return llvm::make_unique<MultiplexConsumer>(std::move(Consumers));
 }
 
-bool GenerateModuleFromModuleMapAction::BeginSourceFileAction(
-    CompilerInstance &CI, StringRef Filename) {
-  return GenerateModuleAction::BeginSourceFileAction(CI, Filename);
-}
-
 std::unique_ptr<raw_pwrite_stream>
 GenerateModuleFromModuleMapAction::CreateOutputFile(CompilerInstance &CI,
                                                     StringRef InFile) {
@@ -194,8 +188,8 @@ GenerateModuleFromModuleMapAction::CreateOutputFile(CompilerInstance &CI,
                              /*CreateMissingDirectories=*/true);
 }
 
-bool GenerateModuleInterfaceAction::BeginSourceFileAction(CompilerInstance &CI,
-                                                          StringRef Filename) {
+bool GenerateModuleInterfaceAction::BeginSourceFileAction(
+    CompilerInstance &CI) {
   if (!CI.getLangOpts().ModulesTS) {
     CI.getDiagnostics().Report(diag::err_module_interface_requires_modules_ts);
     return false;
@@ -203,7 +197,7 @@ bool GenerateModuleInterfaceAction::BeginSourceFileAction(CompilerInstance &CI,
 
   CI.getLangOpts().setCompilingModule(LangOptions::CMK_ModuleInterface);
 
-  return GenerateModuleAction::BeginSourceFileAction(CI, Filename);
+  return GenerateModuleAction::BeginSourceFileAction(CI);
 }
 
 std::unique_ptr<raw_pwrite_stream>
diff --git a/lib/Frontend/Rewrite/CMakeLists.txt b/lib/Frontend/Rewrite/CMakeLists.txt
index 924bf5d5ee284..61a22b5b13b4d 100644
--- a/lib/Frontend/Rewrite/CMakeLists.txt
+++ b/lib/Frontend/Rewrite/CMakeLists.txt
@@ -19,4 +19,5 @@ add_clang_library(clangRewriteFrontend
   clangFrontend
   clangLex
   clangRewrite
+  clangSerialization
   )
diff --git a/lib/Frontend/Rewrite/FrontendActions.cpp b/lib/Frontend/Rewrite/FrontendActions.cpp
index b2dfd29411302..45feffbcb5b54 100644
--- a/lib/Frontend/Rewrite/FrontendActions.cpp
+++ b/lib/Frontend/Rewrite/FrontendActions.cpp
@@ -18,6 +18,11 @@
 #include "clang/Rewrite/Frontend/ASTConsumers.h"
 #include "clang/Rewrite/Frontend/FixItRewriter.h"
 #include "clang/Rewrite/Frontend/Rewriters.h"
+#include "clang/Serialization/ASTReader.h"
+#include "clang/Serialization/Module.h"
+#include "clang/Serialization/ModuleManager.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/CrashRecoveryContext.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
@@ -86,8 +91,7 @@ public:
 };
 } // end anonymous namespace
 
-bool FixItAction::BeginSourceFileAction(CompilerInstance &CI,
-                                        StringRef Filename) {
+bool FixItAction::BeginSourceFileAction(CompilerInstance &CI) {
   const FrontendOptions &FEOpts = getCompilerInstance().getFrontendOpts();
   if (!FEOpts.FixItSuffix.empty()) {
     FixItOpts.reset(new FixItActionSuffixInserter(FEOpts.FixItSuffix,
@@ -190,27 +194,112 @@ void RewriteTestAction::ExecuteAction() {
   DoRewriteTest(CI.getPreprocessor(), OS.get());
 }
 
-void RewriteIncludesAction::ExecuteAction() {
-  CompilerInstance &CI = getCompilerInstance();
-  std::unique_ptr<raw_ostream> OS =
-      CI.createDefaultOutputFile(true, getCurrentFile());
-  if (!OS) return;
+class RewriteIncludesAction::RewriteImportsListener : public ASTReaderListener {
+  CompilerInstance &CI;
+  std::weak_ptr<raw_ostream> Out;
+
+  llvm::DenseSet<const FileEntry*> Rewritten;
+
+public:
+  RewriteImportsListener(CompilerInstance &CI, std::shared_ptr<raw_ostream> Out)
+      : CI(CI), Out(Out) {}
+
+  void visitModuleFile(StringRef Filename,
+                       serialization::ModuleKind Kind) override {
+    auto *File = CI.getFileManager().getFile(Filename);
+    assert(File && "missing file for loaded module?");
+
+    // Only rewrite each module file once.
+    if (!Rewritten.insert(File).second)
+      return;
+
+    serialization::ModuleFile *MF =
+        CI.getModuleManager()->getModuleManager().lookup(File);
+    assert(File && "missing module file for loaded module?");
+
+    // Not interested in PCH / preambles.
+    if (!MF->isModule())
+      return;
+
+    auto OS = Out.lock();
+    assert(OS && "loaded module file after finishing rewrite action?");
+
+    (*OS) << "#pragma clang module build " << MF->ModuleName << "\n";
+
+    // Rewrite the contents of the module in a separate compiler instance.
+    CompilerInstance Instance(CI.getPCHContainerOperations(),
+                              &CI.getPreprocessor().getPCMCache());
+    Instance.setInvocation(
+        std::make_shared<CompilerInvocation>(CI.getInvocation()));
+    Instance.createDiagnostics(
+        new ForwardingDiagnosticConsumer(CI.getDiagnosticClient()),
+        /*ShouldOwnClient=*/true);
+    Instance.getFrontendOpts().Inputs.clear();
+    Instance.getFrontendOpts().Inputs.emplace_back(
+        Filename, InputKind(InputKind::Unknown, InputKind::Precompiled));
+    // Don't recursively rewrite imports. We handle them all at the top level.
+    Instance.getPreprocessorOutputOpts().RewriteImports = false;
+
+    llvm::CrashRecoveryContext().RunSafelyOnThread([&]() {
+      RewriteIncludesAction Action;
+      Action.OutputStream = OS;
+      Instance.ExecuteAction(Action);
+    });
+
+    (*OS) << "#pragma clang module endbuild /*" << MF->ModuleName << "*/\n";
+  }
+};
+
+bool RewriteIncludesAction::BeginSourceFileAction(CompilerInstance &CI) {
+  if (!OutputStream) {
+    OutputStream = CI.createDefaultOutputFile(true, getCurrentFile());
+    if (!OutputStream)
+      return false;
+  }
+
+  auto &OS = *OutputStream;
 
   // If we're preprocessing a module map, start by dumping the contents of the
   // module itself before switching to the input buffer.
   auto &Input = getCurrentInput();
   if (Input.getKind().getFormat() == InputKind::ModuleMap) {
     if (Input.isFile()) {
-      (*OS) << "# 1 \"";
-      OS->write_escaped(Input.getFile());
-      (*OS) << "\"\n";
+      OS << "# 1 \"";
+      OS.write_escaped(Input.getFile());
+      OS << "\"\n";
     }
-    // FIXME: Include additional information here so that we don't need the
-    // original source files to exist on disk.
-    getCurrentModule()->print(*OS);
-    (*OS) << "#pragma clang module contents\n";
+    getCurrentModule()->print(OS);
+    OS << "#pragma clang module contents\n";
+  }
+
+  // If we're rewriting imports, set up a listener to track when we import
+  // module files.
+  if (CI.getPreprocessorOutputOpts().RewriteImports) {
+    CI.createModuleManager();
+    CI.getModuleManager()->addListener(
+        llvm::make_unique<RewriteImportsListener>(CI, OutputStream));
+  }
+
+  return true;
+}
+
+void RewriteIncludesAction::ExecuteAction() {
+  CompilerInstance &CI = getCompilerInstance();
+
+  // If we're rewriting imports, emit the module build output first rather
+  // than switching back and forth (potentially in the middle of a line).
+  if (CI.getPreprocessorOutputOpts().RewriteImports) {
+    std::string Buffer;
+    llvm::raw_string_ostream OS(Buffer);
+
+    RewriteIncludesInInput(CI.getPreprocessor(), &OS,
+                           CI.getPreprocessorOutputOpts());
+
+    (*OutputStream) << OS.str();
+  } else {
+    RewriteIncludesInInput(CI.getPreprocessor(), OutputStream.get(),
+                           CI.getPreprocessorOutputOpts());
   }
 
-  RewriteIncludesInInput(CI.getPreprocessor(), OS.get(),
-                         CI.getPreprocessorOutputOpts());
+  OutputStream.reset();
 }