1 files changed, 6 insertions, 5 deletions

diff --git a/include/llvm/CodeGen/PBQP/Solution.h b/include/llvm/CodeGen/PBQP/Solution.h
index 98ba5a1848..b9f288bbee 100644
--- a/include/llvm/CodeGen/PBQP/Solution.h
+++ b/include/llvm/CodeGen/PBQP/Solution.h
@@ -26,7 +26,8 @@ namespace PBQP {
   class Solution {
   private:
 
-    typedef std::map<Graph::NodeId, unsigned> SelectionsMap;
+    typedef std::map<Graph::ConstNodeItr, unsigned,
+                     NodeItrComparator> SelectionsMap;
     SelectionsMap selections;
 
     unsigned r0Reductions, r1Reductions, r2Reductions, rNReductions;
@@ -72,15 +73,15 @@ namespace PBQP {
     /// \brief Set the selection for a given node.
     /// @param nItr Node iterator.
     /// @param selection Selection for nItr.
-    void setSelection(Graph::NodeId nodeId, unsigned selection) {
-      selections[nodeId] = selection;
+    void setSelection(Graph::NodeItr nItr, unsigned selection) {
+      selections[nItr] = selection;
     }
 
     /// \brief Get a node's selection.
     /// @param nItr Node iterator.
     /// @return The selection for nItr;
-    unsigned getSelection(Graph::NodeId nodeId) const {
-      SelectionsMap::const_iterator sItr = selections.find(nodeId);
+    unsigned getSelection(Graph::ConstNodeItr nItr) const {
+      SelectionsMap::const_iterator sItr = selections.find(nItr);
       assert(sItr != selections.end() && "No selection for node.");
       return sItr->second;
     }