4 files changed, 33 insertions, 29 deletions

diff --git a/include/clang/Basic/SourceManager.h b/include/clang/Basic/SourceManager.h
index d08d5368d1..b1fb78ea56 100644
--- a/include/clang/Basic/SourceManager.h
+++ b/include/clang/Basic/SourceManager.h
@@ -624,6 +624,12 @@ public:
     return MainFileID;
   }
 
+  /// \brief Set the file ID for the main source file.
+  void setMainFileID(FileID FID) {
+    assert(MainFileID.isInvalid() && "MainFileID already set!");
+    MainFileID = FID;
+  }
+
   /// \brief Set the file ID for the precompiled preamble.
   void setPreambleFileID(FileID Preamble) {
     assert(PreambleFileID.isInvalid() && "PreambleFileID already set!");
diff --git a/lib/Frontend/FrontendAction.cpp b/lib/Frontend/FrontendAction.cpp
index 6ef07c447d..7e4ae023d0 100644
--- a/lib/Frontend/FrontendAction.cpp
+++ b/lib/Frontend/FrontendAction.cpp
@@ -310,16 +310,7 @@ void FrontendAction::Execute() {
 
   // Initialize the main file entry. This needs to be delayed until after PCH
   // has loaded.
-  if (isCurrentFileAST()) {
-    // Set the main file ID to an empty file.
-    //
-    // FIXME: We probably shouldn't need this, but for now this is the
-    // simplest way to reuse the logic in ParseAST.
-    const char *EmptyStr = "";
-    llvm::MemoryBuffer *SB =
-      llvm::MemoryBuffer::getMemBuffer(EmptyStr, "<dummy input>");
-    CI.getSourceManager().createMainFileIDForMemBuffer(SB);
-  } else {
+  if (!isCurrentFileAST()) {
     if (!CI.InitializeSourceManager(getCurrentFile()))
       return;
   }
diff --git a/lib/Lex/Preprocessor.cpp b/lib/Lex/Preprocessor.cpp
index cfa9e23e60..5dda96451f 100644
--- a/lib/Lex/Preprocessor.cpp
+++ b/lib/Lex/Preprocessor.cpp
@@ -400,19 +400,23 @@ void Preprocessor::EnterMainSourceFile() {
   assert(NumEnteredSourceFiles == 0 && "Cannot reenter the main file!");
   FileID MainFileID = SourceMgr.getMainFileID();
 
-  // Enter the main file source buffer.
-  EnterSourceFile(MainFileID, 0, SourceLocation());
-
-  // If we've been asked to skip bytes in the main file (e.g., as part of a
-  // precompiled preamble), do so now.
-  if (SkipMainFilePreamble.first > 0)
-    CurLexer->SkipBytes(SkipMainFilePreamble.first, 
-                        SkipMainFilePreamble.second);
+  // If MainFileID is loaded it means we loaded an AST file, no need to enter
+  // a main file.
+  if (!SourceMgr.isLoadedFileID(MainFileID)) {
+    // Enter the main file source buffer.
+    EnterSourceFile(MainFileID, 0, SourceLocation());
   
-  // Tell the header info that the main file was entered.  If the file is later
-  // #imported, it won't be re-entered.
-  if (const FileEntry *FE = SourceMgr.getFileEntryForID(MainFileID))
-    HeaderInfo.IncrementIncludeCount(FE);
+    // If we've been asked to skip bytes in the main file (e.g., as part of a
+    // precompiled preamble), do so now.
+    if (SkipMainFilePreamble.first > 0)
+      CurLexer->SkipBytes(SkipMainFilePreamble.first, 
+                          SkipMainFilePreamble.second);
+    
+    // Tell the header info that the main file was entered.  If the file is later
+    // #imported, it won't be re-entered.
+    if (const FileEntry *FE = SourceMgr.getFileEntryForID(MainFileID))
+      HeaderInfo.IncrementIncludeCount(FE);
+  }
 
   // Preprocess Predefines to populate the initial preprocessor state.
   llvm::MemoryBuffer *SB =
diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index c29033e833..a7d143b030 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -2656,14 +2656,17 @@ ASTReader::ASTReadResult ASTReader::ReadAST(const std::string &FileName,
   if (DeserializationListener)
     DeserializationListener->ReaderInitialized(this);
 
-  // If this AST file is a precompiled preamble, then set the preamble file ID
-  // of the source manager to the file source file from which the preamble was
-  // built.
-  if (Type == MK_Preamble) {
-    if (!OriginalFileID.isInvalid()) {
-      OriginalFileID = FileID::get(ModuleMgr.getPrimaryModule().SLocEntryBaseID
-                                        + OriginalFileID.getOpaqueValue() - 1);
+  if (!OriginalFileID.isInvalid()) {
+    OriginalFileID = FileID::get(ModuleMgr.getPrimaryModule().SLocEntryBaseID
+                                      + OriginalFileID.getOpaqueValue() - 1);
+
+    // If this AST file is a precompiled preamble, then set the preamble file ID
+    // of the source manager to the file source file from which the preamble was
+    // built.
+    if (Type == MK_Preamble) {
       SourceMgr.setPreambleFileID(OriginalFileID);
+    } else if (Type == MK_MainFile) {
+      SourceMgr.setMainFileID(OriginalFileID);
     }
   }