3 files changed, 32 insertions, 11 deletions

diff --git a/lib/Frontend/ASTMerge.cpp b/lib/Frontend/ASTMerge.cpp
index e916e20065..cbaa6a3034 100644
--- a/lib/Frontend/ASTMerge.cpp
+++ b/lib/Frontend/ASTMerge.cpp
@@ -44,6 +44,12 @@ void ASTMergeAction::ExecuteAction() {
     if (!Unit)
       continue;
 
+    // Reset the argument -> string function so that it has the AST
+    // context we want, since the Sema object created by
+    // LoadFromPCHFile will override it.
+    CI.getDiagnostics().SetArgToStringFn(&FormatASTNodeDiagnosticArgument,
+                                         &CI.getASTContext());
+
     ASTImporter Importer(CI.getDiagnostics(),
                          CI.getASTContext(), 
                          CI.getFileManager(),
diff --git a/lib/Frontend/ASTUnit.cpp b/lib/Frontend/ASTUnit.cpp
index bbe2ec5bea..b56a0d83a8 100644
--- a/lib/Frontend/ASTUnit.cpp
+++ b/lib/Frontend/ASTUnit.cpp
@@ -204,7 +204,7 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
   AST->FileMgr.reset(new FileManager);
   AST->SourceMgr.reset(new SourceManager(AST->getDiagnostics()));
   AST->HeaderInfo.reset(new HeaderSearch(AST->getFileManager()));
-
+  
   // If requested, capture diagnostics in the ASTUnit.
   CaptureDroppedDiagnostics Capture(CaptureDiagnostics, AST->getDiagnostics(),
                                     AST->StoredDiagnostics);
@@ -237,7 +237,6 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
   unsigned Counter;
 
   llvm::OwningPtr<PCHReader> Reader;
-  llvm::OwningPtr<ExternalASTSource> Source;
 
   Reader.reset(new PCHReader(AST->getSourceManager(), AST->getFileManager(),
                              AST->getDiagnostics()));
@@ -294,9 +293,18 @@ ASTUnit *ASTUnit::LoadFromPCHFile(const std::string &Filename,
   // Attach the PCH reader to the AST context as an external AST
   // source, so that declarations will be deserialized from the
   // PCH file as needed.
-  Source.reset(Reader.take());
+  PCHReader *ReaderPtr = Reader.get();
+  llvm::OwningPtr<ExternalASTSource> Source(Reader.take());
   Context.setExternalSource(Source);
 
+  // Create an AST consumer, even though it isn't used.
+  AST->Consumer.reset(new ASTConsumer);
+  
+  // Create a semantic analysis object and tell the PCH reader about it.
+  AST->TheSema.reset(new Sema(PP, Context, *AST->Consumer));
+  AST->TheSema->Initialize();
+  ReaderPtr->InitializeSema(*AST->TheSema);
+
   return AST.take();
 }
 
@@ -458,6 +466,7 @@ bool ASTUnit::Parse(llvm::MemoryBuffer *OverrideMainBuffer) {
   // FIXME: Should we retain the previous file manager?
   FileMgr.reset(new FileManager);
   SourceMgr.reset(new SourceManager(getDiagnostics()));
+  TheSema.reset();
   Ctx.reset();
   PP.reset();
   
@@ -513,6 +522,8 @@ bool ASTUnit::Parse(llvm::MemoryBuffer *OverrideMainBuffer) {
   
   // Steal the created target, context, and preprocessor, and take back the
   // source and file managers.
+  TheSema.reset(Clang.takeSema());
+  Consumer.reset(Clang.takeASTConsumer());
   Ctx.reset(Clang.takeASTContext());
   PP.reset(Clang.takePreprocessor());
   Clang.takeSourceManager();
diff --git a/lib/Frontend/PCHReader.cpp b/lib/Frontend/PCHReader.cpp
index e55605b274..137c2d768e 100644
--- a/lib/Frontend/PCHReader.cpp
+++ b/lib/Frontend/PCHReader.cpp
@@ -3135,9 +3135,11 @@ void PCHReader::InitializeSema(Sema &S) {
 
   // Makes sure any declarations that were deserialized "too early"
   // still get added to the identifier's declaration chains.
-  for (unsigned I = 0, N = PreloadedDecls.size(); I != N; ++I) {
-    SemaObj->TUScope->AddDecl(Action::DeclPtrTy::make(PreloadedDecls[I]));
-    SemaObj->IdResolver.AddDecl(PreloadedDecls[I]);
+  if (SemaObj->TUScope) {
+    for (unsigned I = 0, N = PreloadedDecls.size(); I != N; ++I) {
+      SemaObj->TUScope->AddDecl(Action::DeclPtrTy::make(PreloadedDecls[I]));
+      SemaObj->IdResolver.AddDecl(PreloadedDecls[I]);
+    }
   }
   PreloadedDecls.clear();
 
@@ -3332,11 +3334,13 @@ PCHReader::SetGloballyVisibleDecls(IdentifierInfo *II,
   for (unsigned I = 0, N = DeclIDs.size(); I != N; ++I) {
     NamedDecl *D = cast<NamedDecl>(GetDecl(DeclIDs[I]));
     if (SemaObj) {
-      // Introduce this declaration into the translation-unit scope
-      // and add it to the declaration chain for this identifier, so
-      // that (unqualified) name lookup will find it.
-      SemaObj->TUScope->AddDecl(Action::DeclPtrTy::make(D));
-      SemaObj->IdResolver.AddDeclToIdentifierChain(II, D);
+      if (SemaObj->TUScope) {
+        // Introduce this declaration into the translation-unit scope
+        // and add it to the declaration chain for this identifier, so
+        // that (unqualified) name lookup will find it.
+        SemaObj->TUScope->AddDecl(Action::DeclPtrTy::make(D));
+        SemaObj->IdResolver.AddDeclToIdentifierChain(II, D);
+      }
     } else {
       // Queue this declaration so that it will be added to the
       // translation unit scope and identifier's declaration chain