ASTUnit: Explicitly track whether the ASTUnit came from an actual AST or not.

llvm-svn: 90349

2 files changed, 14 insertions, 6 deletions

diff --git a/clang/include/clang/Frontend/ASTUnit.h b/clang/include/clang/Frontend/ASTUnit.h
index 08e8cf5aee4..a0c0e2207cf 100644
--- a/clang/include/clang/Frontend/ASTUnit.h
+++ b/clang/include/clang/Frontend/ASTUnit.h
@@ -53,7 +53,10 @@ class ASTUnit {
   // that come from the AST itself, not from included precompiled headers.
   // FIXME: This is temporary; eventually, CIndex will always do this.
   bool                              OnlyLocalDecls;
-  
+
+  // Track whether the main file was loaded from an AST or not.
+  bool MainFileIsAST;
+
   /// The name of the original source file used to generate this ASTUnit.
   std::string OriginalSourceFile;
 
@@ -64,9 +67,11 @@ class ASTUnit {
   ASTUnit &operator=(const ASTUnit &); // DO NOT IMPLEMENT
 
 public:
-  ASTUnit(DiagnosticClient *diagClient = NULL);
+  ASTUnit(bool MainFileIsAST, DiagnosticClient *diagClient = NULL);
   ~ASTUnit();
 
+  bool isMainFileAST() const { return MainFileIsAST; }
+
   const SourceManager &getSourceManager() const { return SourceMgr; }
         SourceManager &getSourceManager()       { return SourceMgr; }
 
diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index 6cbcb46e526..a3db339b83f 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -34,7 +34,10 @@
 #include "llvm/System/Path.h"
 using namespace clang;
 
-ASTUnit::ASTUnit(DiagnosticClient *diagClient) : tempFile(false) {
+ASTUnit::ASTUnit(bool _MainFileIsAST,
+                 DiagnosticClient *diagClient)
+  : tempFile(false), MainFileIsAST(_MainFileIsAST)
+{
   Diags.setClient(diagClient ? diagClient : new TextDiagnosticBuffer());
 }
 ASTUnit::~ASTUnit() {
@@ -99,7 +102,7 @@ const std::string &ASTUnit::getOriginalSourceFileName() {
 }
 
 const std::string &ASTUnit::getPCHFileName() {
-  assert(Ctx->getExternalSource() && "Not an ASTUnit from a PCH file!");
+  assert(isMainFileAST() && "Not an ASTUnit from a PCH file!");
   return dyn_cast<PCHReader>(Ctx->getExternalSource())->getFileName();
 }
 
@@ -108,7 +111,7 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
                                   DiagnosticClient *diagClient,
                                   bool OnlyLocalDecls,
                                   bool UseBumpAllocator) {
-  llvm::OwningPtr<ASTUnit> AST(new ASTUnit(diagClient));
+  llvm::OwningPtr<ASTUnit> AST(new ASTUnit(true, diagClient));
   AST->OnlyLocalDecls = OnlyLocalDecls;
   AST->HeaderInfo.reset(new HeaderSearch(AST->getFileManager()));
 
@@ -230,7 +233,7 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocation(const CompilerInvocation &CI,
   // Create the AST unit.
   //
   // FIXME: Use the provided diagnostic client.
-  AST.reset(new ASTUnit());
+  AST.reset(new ASTUnit(false));
 
   AST->OnlyLocalDecls = OnlyLocalDecls;
   AST->OriginalSourceFile = Clang.getFrontendOpts().Inputs[0].second;