4 files changed, 93 insertions, 8 deletions

diff --git a/llvm/include/llvm/Analysis/CallGraph.h b/llvm/include/llvm/Analysis/CallGraph.h
index 24fa99f4c18..76d9073799f 100644
--- a/llvm/include/llvm/Analysis/CallGraph.h
+++ b/llvm/include/llvm/Analysis/CallGraph.h
@@ -417,13 +417,24 @@ template <> struct GraphTraits<CallGraphNode *> {
 
 template <> struct GraphTraits<const CallGraphNode *> {
   typedef const CallGraphNode NodeType;
-  typedef NodeType::const_iterator ChildIteratorType;
+
+  typedef CallGraphNode::CallRecord CGNPairTy;
+  typedef std::pointer_to_unary_function<CGNPairTy, const CallGraphNode *>
+      CGNDerefFun;
 
   static NodeType *getEntryNode(const CallGraphNode *CGN) { return CGN; }
+
+  typedef mapped_iterator<NodeType::const_iterator, CGNDerefFun>
+      ChildIteratorType;
+
   static inline ChildIteratorType child_begin(NodeType *N) {
-    return N->begin();
+    return map_iterator(N->begin(), CGNDerefFun(CGNDeref));
+  }
+  static inline ChildIteratorType child_end(NodeType *N) {
+    return map_iterator(N->end(), CGNDerefFun(CGNDeref));
   }
-  static inline ChildIteratorType child_end(NodeType *N) { return N->end(); }
+
+  static const CallGraphNode *CGNDeref(CGNPairTy P) { return P.second; }
 };
 
 template <>
@@ -450,12 +461,22 @@ template <>
 struct GraphTraits<const CallGraph *> : public GraphTraits<
                                             const CallGraphNode *> {
   static NodeType *getEntryNode(const CallGraph *CGN) {
-    return CGN->getExternalCallingNode();
+    return CGN->getExternalCallingNode(); // Start at the external node!
   }
+  typedef std::pair<const Function *, const CallGraphNode *> PairTy;
+  typedef std::pointer_to_unary_function<PairTy, const CallGraphNode &>
+      DerefFun;
+
   // nodes_iterator/begin/end - Allow iteration over all nodes in the graph
-  typedef CallGraph::const_iterator nodes_iterator;
-  static nodes_iterator nodes_begin(const CallGraph *CG) { return CG->begin(); }
-  static nodes_iterator nodes_end(const CallGraph *CG) { return CG->end(); }
+  typedef mapped_iterator<CallGraph::const_iterator, DerefFun> nodes_iterator;
+  static nodes_iterator nodes_begin(const CallGraph *CG) {
+    return map_iterator(CG->begin(), DerefFun(CGdereference));
+  }
+  static nodes_iterator nodes_end(const CallGraph *CG) {
+    return map_iterator(CG->end(), DerefFun(CGdereference));
+  }
+
+  static const CallGraphNode &CGdereference(PairTy P) { return *P.second; }
 };
 
 } // End llvm namespace
diff --git a/llvm/unittests/Analysis/CMakeLists.txt b/llvm/unittests/Analysis/CMakeLists.txt
index 84548609ede..baf0c28a8d9 100644
--- a/llvm/unittests/Analysis/CMakeLists.txt
+++ b/llvm/unittests/Analysis/CMakeLists.txt
@@ -1,4 +1,5 @@
 set(LLVM_LINK_COMPONENTS
+  IPA
   Analysis
   AsmParser
   Core
@@ -6,6 +7,7 @@ set(LLVM_LINK_COMPONENTS
   )
 
 add_llvm_unittest(AnalysisTests
+  CallGraphTest.cpp
   CFGTest.cpp
   LazyCallGraphTest.cpp
   ScalarEvolutionTest.cpp
diff --git a/llvm/unittests/Analysis/CallGraphTest.cpp b/llvm/unittests/Analysis/CallGraphTest.cpp
new file mode 100644
index 00000000000..29e9aa63182
--- /dev/null
+++ b/llvm/unittests/Analysis/CallGraphTest.cpp
@@ -0,0 +1,62 @@
+//=======- CallGraphTest.cpp - Unit tests for the CG analysis -------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Analysis/CallGraph.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+template <typename Ty> void canSpecializeGraphTraitsIterators(Ty *G) {
+  typedef typename GraphTraits<Ty *>::NodeType NodeTy;
+
+  auto I = GraphTraits<Ty *>::nodes_begin(G);
+  auto E = GraphTraits<Ty *>::nodes_end(G);
+  auto X = ++I;
+
+  // Should be able to iterate over all nodes of the graph.
+  static_assert(std::is_same<decltype(*I), NodeTy &>::value,
+                "Node type does not match");
+  static_assert(std::is_same<decltype(*X), NodeTy &>::value,
+                "Node type does not match");
+  static_assert(std::is_same<decltype(*E), NodeTy &>::value,
+                "Node type does not match");
+
+  NodeTy *N = GraphTraits<Ty *>::getEntryNode(G);
+
+  auto S = GraphTraits<NodeTy *>::child_begin(N);
+  auto F = GraphTraits<NodeTy *>::child_end(N);
+  auto Y = ++S;
+
+  // Should be able to iterate over immediate successors of a node.
+  static_assert(std::is_same<decltype(*S), NodeTy *>::value,
+                "Node type does not match");
+  static_assert(std::is_same<decltype(*F), NodeTy *>::value,
+                "Node type does not match");
+  static_assert(std::is_same<decltype(*Y), NodeTy *>::value,
+                "Node type does not match");
+}
+
+TEST(CallGraphTest, GraphTraitsSpecialization) {
+  Module M("", getGlobalContext());
+  CallGraph CG(M);
+
+  canSpecializeGraphTraitsIterators(&CG);
+}
+
+TEST(CallGraphTest, GraphTraitsConstSpecialization) {
+  Module M("", getGlobalContext());
+  CallGraph CG(M);
+
+  canSpecializeGraphTraitsIterators(const_cast<const CallGraph *>(&CG));
+}
+}
diff --git a/llvm/unittests/Analysis/Makefile b/llvm/unittests/Analysis/Makefile
index 527f4525e87..52296e7b3db 100644
--- a/llvm/unittests/Analysis/Makefile
+++ b/llvm/unittests/Analysis/Makefile
@@ -9,7 +9,7 @@
 
 LEVEL = ../..
 TESTNAME = Analysis
-LINK_COMPONENTS := analysis asmparser
+LINK_COMPONENTS := ipa analysis asmparser
 
 include $(LEVEL)/Makefile.config
 include $(LLVM_SRC_ROOT)/unittests/Makefile.unittest