1 files changed, 30 insertions, 21 deletions

diff --git a/lib/Frontend/FrontendAction.cpp b/lib/Frontend/FrontendAction.cpp
index f81a06b31869..704d51509851 100644
--- a/lib/Frontend/FrontendAction.cpp
+++ b/lib/Frontend/FrontendAction.cpp
@@ -200,12 +200,12 @@ FrontendAction::CreateWrappedASTConsumer(CompilerInstance &CI,
 ///
 /// \param CI The compiler instance.
 /// \param InputFile Populated with the filename from the line marker.
-/// \param AddLineNote If \c true, add a line note corresponding to this line
-///        directive. Only use this if the directive will not actually be
-///        visited by the preprocessor.
+/// \param IsModuleMap If \c true, add a line note corresponding to this line
+///        directive. (We need to do this because the directive will not be
+///        visited by the preprocessor.)
 static SourceLocation ReadOriginalFileName(CompilerInstance &CI,
                                            std::string &InputFile,
-                                           bool AddLineNote = false) {
+                                           bool IsModuleMap = false) {
   auto &SourceMgr = CI.getSourceManager();
   auto MainFileID = SourceMgr.getMainFileID();
 
@@ -231,7 +231,7 @@ static SourceLocation ReadOriginalFileName(CompilerInstance &CI,
 
   unsigned LineNo;
   SourceLocation LineNoLoc = T.getLocation();
-  if (AddLineNote) {
+  if (IsModuleMap) {
     llvm::SmallString<16> Buffer;
     if (Lexer::getSpelling(LineNoLoc, Buffer, SourceMgr, CI.getLangOpts())
             .getAsInteger(10, LineNo))
@@ -250,10 +250,10 @@ static SourceLocation ReadOriginalFileName(CompilerInstance &CI,
     return SourceLocation();
   InputFile = Literal.GetString().str();
 
-  if (AddLineNote)
+  if (IsModuleMap)
     CI.getSourceManager().AddLineNote(
         LineNoLoc, LineNo, SourceMgr.getLineTableFilenameID(InputFile), false,
-        false, SrcMgr::C_User);
+        false, SrcMgr::C_User_ModuleMap);
 
   return T.getLocation();
 }
@@ -403,7 +403,7 @@ static bool loadModuleMapForModuleBuild(CompilerInstance &CI, bool IsSystem,
   Offset = 0;
   if (IsPreprocessed) {
     SourceLocation EndOfLineMarker =
-        ReadOriginalFileName(CI, PresumedModuleMapFile, /*AddLineNote*/true);
+        ReadOriginalFileName(CI, PresumedModuleMapFile, /*IsModuleMap*/ true);
     if (EndOfLineMarker.isValid())
       Offset = CI.getSourceManager().getDecomposedLoc(EndOfLineMarker).second;
   }
@@ -536,8 +536,8 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
     ASTDiags->setClient(Diags->getClient(), /*OwnsClient*/false);
 
     std::unique_ptr<ASTUnit> AST = ASTUnit::LoadFromASTFile(
-        InputFile, CI.getPCHContainerReader(), ASTDiags, CI.getFileSystemOpts(),
-        CI.getCodeGenOpts().DebugTypeExtRefs);
+        InputFile, CI.getPCHContainerReader(), ASTUnit::LoadPreprocessorOnly,
+        ASTDiags, CI.getFileSystemOpts(), CI.getCodeGenOpts().DebugTypeExtRefs);
     if (!AST)
       goto failure;
 
@@ -547,27 +547,36 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
     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());
 
+    // Preload all the module files loaded transitively by the AST unit. Also
+    // load all module map files that were parsed as part of building the AST
+    // unit.
+    if (auto ASTReader = AST->getASTReader()) {
+      auto &MM = ASTReader->getModuleManager();
+      auto &PrimaryModule = MM.getPrimaryModule();
+
+      for (ModuleFile &MF : MM)
+        if (&MF != &PrimaryModule)
+          CI.getFrontendOpts().ModuleFiles.push_back(MF.FileName);
+
+      ASTReader->visitTopLevelModuleMaps(PrimaryModule,
+                                         [&](const FileEntry *FE) {
+        CI.getFrontendOpts().ModuleMapFiles.push_back(FE->getName());
+      });
+    }
+
     // 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);
+      assert(ASTModule && "module file does not define its own module");
       Input = FrontendInputFile(ASTModule->PresumedModuleMapFile, Kind);
     } else {
       auto &SM = CI.getSourceManager();
@@ -590,8 +599,8 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
     IntrusiveRefCntPtr<DiagnosticsEngine> Diags(&CI.getDiagnostics());
 
     std::unique_ptr<ASTUnit> AST = ASTUnit::LoadFromASTFile(
-        InputFile, CI.getPCHContainerReader(), Diags, CI.getFileSystemOpts(),
-        CI.getCodeGenOpts().DebugTypeExtRefs);
+        InputFile, CI.getPCHContainerReader(), ASTUnit::LoadEverything, Diags,
+        CI.getFileSystemOpts(), CI.getCodeGenOpts().DebugTypeExtRefs);
 
     if (!AST)
       goto failure;