diff --git a/include/clang/Analysis/CallGraph.h b/include/clang/Analysis/CallGraph.h
index 81926b9..f1b6e64 100644
--- a/include/clang/Analysis/CallGraph.h
+++ b/include/clang/Analysis/CallGraph.h
@@ -48,8 +48,9 @@
   /// \brief Add the given declaration to the call graph.
   void addToCallGraph(Decl *D, bool IsGlobal);
 
-  /// \brief Populate the call graph with the functions in the given DeclContext.
-  void addToCallGraph(DeclContext *DC);
+  /// \brief Populate the call graph with the functions in the given translation
+  /// unit.
+  void addToCallGraph(TranslationUnitDecl *TU);
 
   /// \brief Lookup the node for the given declaration.
   CallGraphNode *getNode(const Decl *) const;
diff --git a/lib/Analysis/CallGraph.cpp b/lib/Analysis/CallGraph.cpp
index cc6e62c..1e9aec2 100644
--- a/lib/Analysis/CallGraph.cpp
+++ b/lib/Analysis/CallGraph.cpp
@@ -14,6 +14,7 @@
 
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/StmtVisitor.h"
 
 #include "llvm/Support/GraphWriter.h"
@@ -71,6 +72,35 @@
         static_cast<CGBuilder*>(this)->Visit(*I);
   }
 };
+
+/// A helper class which walks the AST declarations.
+// TODO: We might want to specialize the visitor to shrink the call graph.
+// For example, we might not want to include the inline methods from header
+// files.
+class CGDeclVisitor : public RecursiveASTVisitor<CGDeclVisitor> {
+  CallGraph *CG;
+
+public:
+  CGDeclVisitor(CallGraph * InCG) : CG(InCG) {}
+
+  bool VisitFunctionDecl(FunctionDecl *FD) {
+    // We skip function template definitions, as their semantics is
+    // only determined when they are instantiated.
+    if (includeInGraph(FD))
+      // If this function has external linkage, anything could call it.
+      // Note, we are not precise here. For example, the function could have
+      // its address taken.
+      CG->addToCallGraph(FD, FD->isGlobal());
+    return true;
+  }
+
+  bool VisitObjCMethodDecl(ObjCMethodDecl *MD) {
+    if (includeInGraph(MD))
+      CG->addToCallGraph(MD, true);
+    return true;
+  }
+};
+
 } // end anonymous namespace
 
 CallGraph::CallGraph() {
@@ -98,42 +128,8 @@
   builder.Visit(D->getBody());
 }
 
-void CallGraph::addToCallGraph(DeclContext *DC) {
-  for (DeclContext::decl_iterator I = DC->decls_begin(), E = DC->decls_end();
-       I != E; ++I) {
-    Decl *D = *I;
-    switch (D->getKind()) {
-    case Decl::CXXConstructor:
-    case Decl::CXXDestructor:
-    case Decl::CXXConversion:
-    case Decl::CXXMethod:
-    case Decl::Function: {
-      FunctionDecl *FD = cast<FunctionDecl>(D);
-      // We skip function template definitions, as their semantics is
-      // only determined when they are instantiated.
-      if (includeInGraph(FD))
-        // If this function has external linkage, anything could call it.
-        // Note, we are not precise here. For example, the function could have
-        // its address taken.
-        addToCallGraph(FD, FD->isGlobal());
-      break;
-    }
-
-    case Decl::ObjCCategoryImpl:
-    case Decl::ObjCImplementation: {
-      ObjCImplDecl *ID = cast<ObjCImplDecl>(D);
-      for (ObjCContainerDecl::method_iterator MI = ID->meth_begin(),
-          ME = ID->meth_end(); MI != ME; ++MI) {
-        if (includeInGraph(*MI))
-          addToCallGraph(*MI, true);
-      }
-      break;
-    }
-
-    default:
-      break;
-    }
-  }
+void CallGraph::addToCallGraph(TranslationUnitDecl *TU) {
+  CGDeclVisitor(this).TraverseDecl(TU);
 }
 
 CallGraphNode *CallGraph::getNode(const Decl *F) const {
diff --git a/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp b/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp
index 93a3ab4..a97ff4a 100644
--- a/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp
+++ b/lib/StaticAnalyzer/Frontend/AnalysisConsumer.cpp
@@ -206,9 +206,9 @@
 
   virtual void HandleTranslationUnit(ASTContext &C);
 
-  /// \brief Build the call graph for the context and use it to define the order
+  /// \brief Build the call graph for the TU and use it to define the order
   /// in which the functions should be visited.
-  void HandleDeclContextGallGraph(ASTContext &C, DeclContext *dc);
+  void HandleDeclsGallGraph(TranslationUnitDecl *TU);
 
   /// \brief Run analyzes(syntax or path sensitive) on the given function.
   /// \param Mode - determines if we are requesting syntax only or path
@@ -261,12 +261,11 @@
 //===----------------------------------------------------------------------===//
 llvm::Timer* AnalysisConsumer::TUTotalTimer = 0;
 
-void AnalysisConsumer::HandleDeclContextGallGraph(ASTContext &C,
-                                                  DeclContext *dc) {
+void AnalysisConsumer::HandleDeclsGallGraph(TranslationUnitDecl *TU) {
   // Otherwise, use the Callgraph to derive the order.
   // Build the Call Graph.
   CallGraph CG;
-  CG.addToCallGraph(dc);
+  CG.addToCallGraph(TU);
 
   // Find the top level nodes - children of root + the unreachable (parentless)
   // nodes.
@@ -334,7 +333,7 @@
     TraverseDecl(TU);
 
     if (Mgr->shouldInlineCall())
-      HandleDeclContextGallGraph(C, TU);
+      HandleDeclsGallGraph(TU);
 
     // After all decls handled, run checkers on the entire TranslationUnit.
     checkerMgr->runCheckersOnEndOfTranslationUnit(TU, *Mgr, BR);