1 files changed, 20 insertions, 11 deletions
diff --git a/lib/Frontend/FrontendAction.cpp b/lib/Frontend/FrontendAction.cpp
index e3d8b85..42da44c 100644
--- a/lib/Frontend/FrontendAction.cpp
+++ b/lib/Frontend/FrontendAction.cpp
@@ -20,6 +20,7 @@
 #include "clang/Frontend/MultiplexConsumer.h"
 #include "clang/Parse/ParseAST.h"
 #include "clang/Serialization/ASTDeserializationListener.h"
+#include "clang/Serialization/ChainedIncludesSource.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Timer.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -209,8 +210,16 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
 
     CI.getASTContext().setASTMutationListener(Consumer->GetASTMutationListener());
 
-    /// Use PCH?
-    if (!CI.getPreprocessorOpts().ImplicitPCHInclude.empty()) {
+    if (!CI.getPreprocessorOpts().ChainedIncludes.empty()) {
+      // Convert headers to PCH and chain them.
+      llvm::OwningPtr<ExternalASTSource> source;
+      source.reset(ChainedIncludesSource::create(CI));
+      if (!source)
+        goto failure;
+      CI.getASTContext().setExternalSource(source);
+
+    } else if (!CI.getPreprocessorOpts().ImplicitPCHInclude.empty()) {
+      // Use PCH.
       assert(hasPCHSupport() && "This action does not have PCH support!");
       ASTDeserializationListener *DeserialListener
           = CI.getInvocation().getFrontendOpts().ChainedPCH ?
@@ -249,10 +258,10 @@ bool FrontendAction::BeginSourceFile(CompilerInstance &CI,
   // matching EndSourceFile().
   failure:
   if (isCurrentFileAST()) {
-    CI.takeASTContext();
-    CI.takePreprocessor();
-    CI.takeSourceManager();
-    CI.takeFileManager();
+    CI.setASTContext(0);
+    CI.setPreprocessor(0);
+    CI.setSourceManager(0);
+    CI.setFileManager(0);
   }
 
   CI.getDiagnosticClient().EndSourceFile();
@@ -304,7 +313,7 @@ void FrontendAction::EndSourceFile() {
     CI.takeASTConsumer();
     if (!isCurrentFileAST()) {
       CI.takeSema();
-      CI.takeASTContext();
+      CI.resetAndLeakASTContext();
     }
   } else {
     if (!isCurrentFileAST()) {
@@ -333,10 +342,10 @@ void FrontendAction::EndSourceFile() {
 
   if (isCurrentFileAST()) {
     CI.takeSema();
-    CI.takeASTContext();
-    CI.takePreprocessor();
-    CI.takeSourceManager();
-    CI.takeFileManager();
+    CI.resetAndLeakASTContext();
+    CI.resetAndLeakPreprocessor();
+    CI.resetAndLeakSourceManager();
+    CI.resetAndLeakFileManager();
   }
 
   setCompilerInstance(0);