diff --git a/clang/include/clang/Analysis/CallGraph.h b/clang/include/clang/Analysis/CallGraph.h
index fabeea38d597..5edfe6fea8db 100644
--- a/clang/include/clang/Analysis/CallGraph.h
+++ b/clang/include/clang/Analysis/CallGraph.h
@@ -17,7 +17,6 @@
 #include "clang/Index/ASTLocation.h"
 #include "clang/Index/Entity.h"
 #include "clang/Index/Program.h"
-#include "clang/Frontend/ASTUnit.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/GraphTraits.h"
 #include "llvm/ADT/STLExtras.h"
@@ -87,7 +86,7 @@ public:
 
   CallGraphNode *getExternalCallingNode() { return ExternalCallingNode; }
 
-  void addTU(ASTUnit &AST);
+  void addTU(ASTContext &AST);
 
   idx::Program &getProgram() { return Prog; }
 
diff --git a/clang/lib/Analysis/CallGraph.cpp b/clang/lib/Analysis/CallGraph.cpp
index ae8845db63ae..17dc0685f82f 100644
--- a/clang/lib/Analysis/CallGraph.cpp
+++ b/clang/lib/Analysis/CallGraph.cpp
@@ -68,10 +68,8 @@ CallGraph::~CallGraph() {
   }
 }
 
-void CallGraph::addTU(ASTUnit &AST) {
-  ASTContext &Ctx = AST.getASTContext();
+void CallGraph::addTU(ASTContext& Ctx) {
   DeclContext *DC = Ctx.getTranslationUnitDecl();
-
   for (DeclContext::decl_iterator I = DC->decls_begin(), E = DC->decls_end();
        I != E; ++I) {
 
diff --git a/clang/tools/wpa/clang-wpa.cpp b/clang/tools/wpa/clang-wpa.cpp
index fa2326dc2ba4..346634b6f617 100644
--- a/clang/tools/wpa/clang-wpa.cpp
+++ b/clang/tools/wpa/clang-wpa.cpp
@@ -13,7 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/Analysis/CallGraph.h"
-
+#include "clang/Frontend/ASTUnit.h"
 #include "clang/Basic/FileManager.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Frontend/TextDiagnosticBuffer.h"
@@ -56,7 +56,7 @@ int main(int argc, char **argv) {
   CG.reset(new CallGraph());
 
   for (unsigned i = 0, e = ASTUnits.size(); i != e; ++i)
-    CG->addTU(*ASTUnits[i]);
+    CG->addTU(ASTUnits[i]->getASTContext());
 
   CG->ViewCallGraph();
 }