1 files changed, 37 insertions, 23 deletions

diff --git a/lib/CodeGen/RegAllocBasic.cpp b/lib/CodeGen/RegAllocBasic.cpp
index 045c8db..6923908 100644
--- a/lib/CodeGen/RegAllocBasic.cpp
+++ b/lib/CodeGen/RegAllocBasic.cpp
@@ -45,6 +45,7 @@
 #include "llvm/Support/Timer.h"
 
 #include <cstdlib>
+#include <queue>
 
 using namespace llvm;
 
@@ -65,6 +66,14 @@ const char *RegAllocBase::TimerGroupName = "Register Allocation";
 bool RegAllocBase::VerifyEnabled = false;
 
 namespace {
+  struct CompSpillWeight {
+    bool operator()(LiveInterval *A, LiveInterval *B) const {
+      return A->weight < B->weight;
+    }
+  };
+}
+
+namespace {
 /// RABasic provides a minimal implementation of the basic register allocation
 /// algorithm. It prioritizes live virtual registers by spill weight and spills
 /// whenever a register is unavailable. This is not practical in production but
@@ -82,7 +91,8 @@ class RABasic : public MachineFunctionPass, public RegAllocBase
 
   // state
   std::auto_ptr<Spiller> SpillerInstance;
-
+  std::priority_queue<LiveInterval*, std::vector<LiveInterval*>,
+                      CompSpillWeight> Queue;
 public:
   RABasic();
 
@@ -100,6 +110,18 @@ public:
 
   virtual float getPriority(LiveInterval *LI) { return LI->weight; }
 
+  virtual void enqueue(LiveInterval *LI) {
+    Queue.push(LI);
+  }
+
+  virtual LiveInterval *dequeue() {
+    if (Queue.empty())
+      return 0;
+    LiveInterval *LI = Queue.top();
+    Queue.pop();
+    return LI;
+  }
+
   virtual unsigned selectOrSplit(LiveInterval &VirtReg,
                                  SmallVectorImpl<LiveInterval*> &SplitVRegs);
 
@@ -227,18 +249,17 @@ void RegAllocBase::releaseMemory() {
   PhysReg2LiveUnion.clear();
 }
 
-// Visit all the live virtual registers. If they are already assigned to a
-// physical register, unify them with the corresponding LiveIntervalUnion,
-// otherwise push them on the priority queue for later assignment.
-void RegAllocBase::
-seedLiveVirtRegs(std::priority_queue<std::pair<float, unsigned> > &VirtRegQ) {
+// Visit all the live registers. If they are already assigned to a physical
+// register, unify them with the corresponding LiveIntervalUnion, otherwise push
+// them on the priority queue for later assignment.
+void RegAllocBase::seedLiveRegs() {
   for (LiveIntervals::iterator I = LIS->begin(), E = LIS->end(); I != E; ++I) {
     unsigned RegNum = I->first;
     LiveInterval &VirtReg = *I->second;
     if (TargetRegisterInfo::isPhysicalRegister(RegNum))
       PhysReg2LiveUnion[RegNum].unify(VirtReg);
     else
-      VirtRegQ.push(std::make_pair(getPriority(&VirtReg), RegNum));
+      enqueue(&VirtReg);
   }
 }
 
@@ -263,38 +284,31 @@ void RegAllocBase::unassign(LiveInterval &VirtReg, unsigned PhysReg) {
 // Top-level driver to manage the queue of unassigned VirtRegs and call the
 // selectOrSplit implementation.
 void RegAllocBase::allocatePhysRegs() {
-
-  // Push each vreg onto a queue or "precolor" by adding it to a physreg union.
-  std::priority_queue<std::pair<float, unsigned> > VirtRegQ;
-  seedLiveVirtRegs(VirtRegQ);
+  seedLiveRegs();
 
   // Continue assigning vregs one at a time to available physical registers.
-  while (!VirtRegQ.empty()) {
-    // Pop the highest priority vreg.
-    LiveInterval &VirtReg = LIS->getInterval(VirtRegQ.top().second);
-    VirtRegQ.pop();
-
+  while (LiveInterval *VirtReg = dequeue()) {
     // selectOrSplit requests the allocator to return an available physical
     // register if possible and populate a list of new live intervals that
     // result from splitting.
-    DEBUG(dbgs() << "\nselectOrSplit " << MRI->getRegClass(VirtReg.reg)->getName()
-                 << ':' << VirtReg << '\n');
+    DEBUG(dbgs() << "\nselectOrSplit "
+                 << MRI->getRegClass(VirtReg->reg)->getName()
+                 << ':' << *VirtReg << '\n');
     typedef SmallVector<LiveInterval*, 4> VirtRegVec;
     VirtRegVec SplitVRegs;
-    unsigned AvailablePhysReg = selectOrSplit(VirtReg, SplitVRegs);
+    unsigned AvailablePhysReg = selectOrSplit(*VirtReg, SplitVRegs);
 
     if (AvailablePhysReg)
-      assign(VirtReg, AvailablePhysReg);
+      assign(*VirtReg, AvailablePhysReg);
 
     for (VirtRegVec::iterator I = SplitVRegs.begin(), E = SplitVRegs.end();
          I != E; ++I) {
-      LiveInterval* SplitVirtReg = *I;
+      LiveInterval *SplitVirtReg = *I;
       if (SplitVirtReg->empty()) continue;
       DEBUG(dbgs() << "queuing new interval: " << *SplitVirtReg << "\n");
       assert(TargetRegisterInfo::isVirtualRegister(SplitVirtReg->reg) &&
              "expect split value in virtual register");
-      VirtRegQ.push(std::make_pair(getPriority(SplitVirtReg),
-                                   SplitVirtReg->reg));
+      enqueue(SplitVirtReg);
       ++NumNewQueued;
     }
   }