1 files changed, 23 insertions, 12 deletions
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index d3475b5..c0415bf 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -21,11 +21,20 @@
 #include "clang/Basic/TargetInfo.h"
 #include "clang/Basic/Diagnostic.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/System/Path.h"
 
 using namespace clang;
 
-ASTUnit::ASTUnit(Diagnostic &_Diags) : Diags(_Diags) { }
-ASTUnit::~ASTUnit() { }
+ASTUnit::ASTUnit(DiagnosticClient *diagClient) : tempFile(false) {      
+  Diags.setClient(diagClient ? diagClient : new TextDiagnosticBuffer());
+}
+ASTUnit::~ASTUnit() { 
+  if (tempFile)
+    llvm::sys::Path(getPCHFileName()).eraseFromDisk();
+  
+  //  The ASTUnit object owns the DiagnosticClient.
+  delete Diags.getClient();
+}
 
 namespace {
 
@@ -80,17 +89,18 @@ const std::string &ASTUnit::getOriginalSourceFileName() {
   return dyn_cast<PCHReader>(Ctx->getExternalSource())->getOriginalSourceFile();
 }
 
-FileManager &ASTUnit::getFileManager() {
-  return HeaderInfo->getFileMgr();
+const std::string &ASTUnit::getPCHFileName() {
+  return dyn_cast<PCHReader>(Ctx->getExternalSource())->getFileName();
 }
 
 ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
-                                  Diagnostic &Diags,
-                                  FileManager &FileMgr,
-                                  std::string *ErrMsg) {
-  llvm::OwningPtr<ASTUnit> AST(new ASTUnit(Diags));
-
-  AST->HeaderInfo.reset(new HeaderSearch(FileMgr));
+                                  std::string *ErrMsg,
+                                  DiagnosticClient *diagClient,
+                                  bool OnlyLocalDecls,
+                                  bool UseBumpAllocator) {
+  llvm::OwningPtr<ASTUnit> AST(new ASTUnit(diagClient));
+  AST->OnlyLocalDecls = OnlyLocalDecls;
+  AST->HeaderInfo.reset(new HeaderSearch(AST->getFileManager()));
 
   // Gather Info for preprocessor construction later on.
 
@@ -103,7 +113,8 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
   llvm::OwningPtr<PCHReader> Reader;
   llvm::OwningPtr<ExternalASTSource> Source;
 
-  Reader.reset(new PCHReader(AST->getSourceManager(), FileMgr, AST->Diags));
+  Reader.reset(new PCHReader(AST->getSourceManager(), AST->getFileManager(),
+                             AST->Diags));
   Reader->setListener(new PCHInfoCollector(LangInfo, HeaderInfo, TargetTriple,
                                            Predefines, Counter));
 
@@ -138,7 +149,7 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
                                 PP.getIdentifierTable(),
                                 PP.getSelectorTable(),
                                 PP.getBuiltinInfo(),
-                                /* FreeMemory = */ true,
+                                /* FreeMemory = */ !UseBumpAllocator,
                                 /* size_reserve = */0));
   ASTContext &Context = *AST->Ctx.get();