1 files changed, 16 insertions, 2 deletions
diff --git a/lib/Lex/HeaderSearch.cpp b/lib/Lex/HeaderSearch.cpp
index d4d14e8745..9e4b68232a 100644
--- a/lib/Lex/HeaderSearch.cpp
+++ b/lib/Lex/HeaderSearch.cpp
@@ -47,7 +47,7 @@ HeaderSearch::HeaderSearch(IntrusiveRefCntPtr<HeaderSearchOptions> HSOpts,
                            const LangOptions &LangOpts, 
                            const TargetInfo *Target)
   : HSOpts(HSOpts), FileMgr(FM), FrameworkMap(64),
-    ModMap(FileMgr, *Diags.getClient(), LangOpts, Target)
+    ModMap(FileMgr, *Diags.getClient(), LangOpts, Target, *this)
 {
   AngledDirIdx = 0;
   SystemDirIdx = 0;
@@ -806,6 +806,7 @@ static void mergeHeaderFileInfo(HeaderFileInfo &HFI,
                                 const HeaderFileInfo &OtherHFI) {
   HFI.isImport |= OtherHFI.isImport;
   HFI.isPragmaOnce |= OtherHFI.isPragmaOnce;
+  HFI.isModuleHeader |= OtherHFI.isModuleHeader;
   HFI.NumIncludes += OtherHFI.NumIncludes;
   
   if (!HFI.ControllingMacro && !HFI.ControllingMacroID) {
@@ -851,6 +852,14 @@ bool HeaderSearch::isFileMultipleIncludeGuarded(const FileEntry *File) {
       HFI.ControllingMacro || HFI.ControllingMacroID;
 }
 
+void HeaderSearch::MarkFileModuleHeader(const FileEntry *FE) {
+  if (FE->getUID() >= FileInfo.size())
+    FileInfo.resize(FE->getUID()+1);
+
+  HeaderFileInfo &HFI = FileInfo[FE->getUID()];
+  HFI.isModuleHeader = true;
+}
+
 void HeaderSearch::setHeaderFileInfoForUID(HeaderFileInfo HFI, unsigned UID) {
   if (UID >= FileInfo.size())
     FileInfo.resize(UID+1);
@@ -948,7 +957,12 @@ bool HeaderSearch::hasModuleMap(StringRef FileName,
   } while (true);
 }
 
-Module *HeaderSearch::findModuleForHeader(const FileEntry *File) {
+Module *HeaderSearch::findModuleForHeader(const FileEntry *File) const {
+  if (ExternalSource) {
+    // Make sure the external source has handled header info about this file,
+    // which includes whether the file is part of a module.
+    (void)getFileInfo(File);
+  }
   if (Module *Mod = ModMap.findModuleForHeader(File))
     return Mod;