3 files changed, 34 insertions, 9 deletions
diff --git a/lib/CodeGen/InstrSched/SchedGraph.h b/lib/CodeGen/InstrSched/SchedGraph.h
index 5c4a600790..ef3b4df862 100644
--- a/lib/CodeGen/InstrSched/SchedGraph.h
+++ b/lib/CodeGen/InstrSched/SchedGraph.h
@@ -19,13 +19,14 @@
 #ifndef LLVM_CODEGEN_SCHEDGRAPH_H
 #define LLVM_CODEGEN_SCHEDGRAPH_H
 
-#include "llvm/CFGdecls.h"			// just for graph iterators
 #include "llvm/Support/NonCopyable.h"
 #include "llvm/Support/HashExtras.h"
+#include "llvm/Support/GraphTraits.h"
 #include <hash_map>
 
 class Value;
 class Instruction;
+class TerminatorInst;
 class BasicBlock;
 class Method;
 class TargetMachine;
@@ -480,13 +481,36 @@ inline sg_succ_const_iterator succ_end(  const SchedGraphNode *N) {
   return sg_succ_const_iterator(N->endOutEdges());
 }
 
-// 
-// po_iterator
-// po_const_iterator
+// Provide specializations of GraphTraits to be able to use graph iterators on
+// the scheduling graph!
 //
-typedef cfg::POIterator<SchedGraphNode, sg_succ_iterator> sg_po_iterator;
-typedef cfg::POIterator<const SchedGraphNode, 
-		        sg_succ_const_iterator> sg_po_const_iterator;
+template <> struct GraphTraits<SchedGraph*> {
+  typedef SchedGraphNode NodeType;
+  typedef sg_succ_iterator ChildIteratorType;
+
+  static inline NodeType *getEntryNode(SchedGraph *SG) { return SG->getRoot(); }
+  static inline ChildIteratorType child_begin(NodeType *N) { 
+    return succ_begin(N); 
+  }
+  static inline ChildIteratorType child_end(NodeType *N) { 
+    return succ_end(N);
+  }
+};
+
+template <> struct GraphTraits<const SchedGraph*> {
+  typedef const SchedGraphNode NodeType;
+  typedef sg_succ_const_iterator ChildIteratorType;
+
+  static inline NodeType *getEntryNode(const SchedGraph *SG) {
+    return SG->getRoot();
+  }
+  static inline ChildIteratorType child_begin(NodeType *N) { 
+    return succ_begin(N); 
+  }
+  static inline ChildIteratorType child_end(NodeType *N) { 
+    return succ_end(N);
+  }
+};
 
 
 //************************ External Functions *****************************/
diff --git a/lib/CodeGen/InstrSched/SchedPriorities.cpp b/lib/CodeGen/InstrSched/SchedPriorities.cpp
index 9e2053bb7c..7840a25098 100644
--- a/lib/CodeGen/InstrSched/SchedPriorities.cpp
+++ b/lib/CodeGen/InstrSched/SchedPriorities.cpp
@@ -19,6 +19,7 @@
 //**************************************************************************/
 
 #include "SchedPriorities.h"
+#include "llvm/Support/PostOrderIterator.h"
 
 
 SchedPriorities::SchedPriorities(const Method* method,
@@ -50,8 +51,7 @@ SchedPriorities::initialize()
 void
 SchedPriorities::computeDelays(const SchedGraph* graph)
 {
-  sg_po_const_iterator poIter = sg_po_const_iterator::begin(graph->getRoot());
-  sg_po_const_iterator poEnd  = sg_po_const_iterator::end(  graph->getRoot());
+  po_iterator<const SchedGraph*> poIter = po_begin(graph), poEnd =po_end(graph);
   for ( ; poIter != poEnd; ++poIter)
     {
       const SchedGraphNode* node = *poIter;
diff --git a/lib/CodeGen/InstrSched/SchedPriorities.h b/lib/CodeGen/InstrSched/SchedPriorities.h
index 909f4ed0eb..81a2e6a053 100644
--- a/lib/CodeGen/InstrSched/SchedPriorities.h
+++ b/lib/CodeGen/InstrSched/SchedPriorities.h
@@ -25,6 +25,7 @@
 #include "llvm/CodeGen/InstrScheduling.h"
 #include "llvm/Analysis/LiveVar/MethodLiveVarInfo.h"
 #include "llvm/Target/MachineSchedInfo.h"
+#include <list>
 
 class Method;
 class MachineInstr;