Fix ASTUnit::getOriginalSourceFileName() when using ASTUnit's derived from

source files.

llvm-svn: 90311

2 files changed, 9 insertions, 1 deletions

diff --git a/clang/include/clang/Frontend/ASTUnit.h b/clang/include/clang/Frontend/ASTUnit.h
index c55638a2370..31d0b26c0d2 100644
--- a/clang/include/clang/Frontend/ASTUnit.h
+++ b/clang/include/clang/Frontend/ASTUnit.h
@@ -54,6 +54,9 @@ class ASTUnit {
   // FIXME: This is temporary; eventually, CIndex will always do this.
   bool                              OnlyLocalDecls;
   
+  /// The name of the original source file used to generate this ASTUnit.
+  std::string OriginalSourceFile;
+
   // Critical optimization when using clang_getCursor().
   ASTLocation LastLoc;
   
diff --git a/clang/lib/Frontend/ASTUnit.cpp b/clang/lib/Frontend/ASTUnit.cpp
index ac213380588..87d0d7d1a46 100644
--- a/clang/lib/Frontend/ASTUnit.cpp
+++ b/clang/lib/Frontend/ASTUnit.cpp
@@ -95,10 +95,11 @@ public:
 } // anonymous namespace
 
 const std::string &ASTUnit::getOriginalSourceFileName() {
-  return dyn_cast<PCHReader>(Ctx->getExternalSource())->getOriginalSourceFile();
+  return OriginalSourceFile;
 }
 
 const std::string &ASTUnit::getPCHFileName() {
+  assert(Ctx->getExternalSource() && "Not an ASTUnit from a PCH file!");
   return dyn_cast<PCHReader>(Ctx->getExternalSource())->getFileName();
 }
 
@@ -138,6 +139,8 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
     return NULL;
   }
 
+  AST->OriginalSourceFile = Reader->getOriginalSourceFile();
+
   // PCH loaded successfully. Now create the preprocessor.
 
   // Get information about the target being compiled for.
@@ -230,6 +233,8 @@ ASTUnit *ASTUnit::LoadFromCompilerInvocation(const CompilerInvocation &CI,
   // FIXME: Use the provided diagnostic client.
   AST.reset(new ASTUnit());
 
+  AST->OriginalSourceFile = Clang.getFrontendOpts().Inputs[0].second;
+
   // Create a file manager object to provide access to and cache the filesystem.
   Clang.setFileManager(&AST->getFileManager());