1 files changed, 58 insertions, 31 deletions

diff --git a/include/llvm/CodeGen/RegisterPressure.h b/include/llvm/CodeGen/RegisterPressure.h
index 9d8843d1d74a5..987634fb36c39 100644
--- a/include/llvm/CodeGen/RegisterPressure.h
+++ b/include/llvm/CodeGen/RegisterPressure.h
@@ -125,11 +125,13 @@ class PressureDiff {
   enum { MaxPSets = 16 };
 
   PressureChange PressureChanges[MaxPSets];
-public:
+
   typedef PressureChange* iterator;
+  iterator nonconst_begin() { return &PressureChanges[0]; }
+  iterator nonconst_end() { return &PressureChanges[MaxPSets]; }
+
+public:
   typedef const PressureChange* const_iterator;
-  iterator begin() { return &PressureChanges[0]; }
-  iterator end() { return &PressureChanges[MaxPSets]; }
   const_iterator begin() const { return &PressureChanges[0]; }
   const_iterator end() const { return &PressureChanges[MaxPSets]; }
 
@@ -191,30 +193,56 @@ struct RegPressureDelta {
   }
 };
 
-/// \brief A set of live virtual registers and physical register units.
+/// A set of live virtual registers and physical register units.
 ///
-/// Virtual and physical register numbers require separate sparse sets, but most
-/// of the RegisterPressureTracker handles them uniformly.
-struct LiveRegSet {
-  SparseSet<unsigned> PhysRegs;
-  SparseSet<unsigned, VirtReg2IndexFunctor> VirtRegs;
+/// This is a wrapper around a SparseSet which deals with mapping register unit
+/// and virtual register indexes to an index usable by the sparse set.
+class LiveRegSet {
+private:
+  SparseSet<unsigned> Regs;
+  unsigned NumRegUnits;
 
-  bool contains(unsigned Reg) const {
+  unsigned getSparseIndexFromReg(unsigned Reg) const {
     if (TargetRegisterInfo::isVirtualRegister(Reg))
-      return VirtRegs.count(Reg);
-    return PhysRegs.count(Reg);
+      return TargetRegisterInfo::virtReg2Index(Reg) + NumRegUnits;
+    assert(Reg < NumRegUnits);
+    return Reg;
+  }
+  unsigned getRegFromSparseIndex(unsigned SparseIndex) const {
+    if (SparseIndex >= NumRegUnits)
+      return TargetRegisterInfo::index2VirtReg(SparseIndex-NumRegUnits);
+    return SparseIndex;
+  }
+
+public:
+  void clear();
+  void init(const MachineRegisterInfo &MRI);
+
+  bool contains(unsigned Reg) const {
+    unsigned SparseIndex = getSparseIndexFromReg(Reg);
+    return Regs.count(SparseIndex);
   }
 
   bool insert(unsigned Reg) {
-    if (TargetRegisterInfo::isVirtualRegister(Reg))
-      return VirtRegs.insert(Reg).second;
-    return PhysRegs.insert(Reg).second;
+    unsigned SparseIndex = getSparseIndexFromReg(Reg);
+    return Regs.insert(SparseIndex).second;
   }
 
   bool erase(unsigned Reg) {
-    if (TargetRegisterInfo::isVirtualRegister(Reg))
-      return VirtRegs.erase(Reg);
-    return PhysRegs.erase(Reg);
+    unsigned SparseIndex = getSparseIndexFromReg(Reg);
+    return Regs.erase(SparseIndex);
+  }
+
+  size_t size() const {
+    return Regs.size();
+  }
+
+  template<typename ContainerT>
+  void appendTo(ContainerT &To) const {
+    for (unsigned I : Regs) {
+      unsigned Reg = getRegFromSparseIndex(I);
+      To.push_back(Reg);
+    }
   }
 };
 
@@ -300,16 +328,12 @@ public:
   // position changes while pressure does not.
   void setPos(MachineBasicBlock::const_iterator Pos) { CurrPos = Pos; }
 
-  /// \brief Get the SlotIndex for the first nondebug instruction including or
-  /// after the current position.
-  SlotIndex getCurrSlot() const;
-
   /// Recede across the previous instruction.
-  bool recede(SmallVectorImpl<unsigned> *LiveUses = nullptr,
+  void recede(SmallVectorImpl<unsigned> *LiveUses = nullptr,
               PressureDiff *PDiff = nullptr);
 
   /// Advance across the current instruction.
-  bool advance();
+  void advance();
 
   /// Finalize the region boundaries and recored live ins and live outs.
   void closeRegion();
@@ -326,17 +350,15 @@ public:
   ArrayRef<unsigned> getLiveThru() const { return LiveThruPressure; }
 
   /// Get the resulting register pressure over the traversed region.
-  /// This result is complete if either advance() or recede() has returned true,
-  /// or if closeRegion() was explicitly invoked.
+  /// This result is complete if closeRegion() was explicitly invoked.
   RegisterPressure &getPressure() { return P; }
   const RegisterPressure &getPressure() const { return P; }
 
   /// Get the register set pressure at the current position, which may be less
   /// than the pressure across the traversed region.
-  std::vector<unsigned> &getRegSetPressureAtPos() { return CurrSetPressure; }
-
-  void discoverLiveOut(unsigned Reg);
-  void discoverLiveIn(unsigned Reg);
+  const std::vector<unsigned> &getRegSetPressureAtPos() const {
+    return CurrSetPressure;
+  }
 
   bool isTopClosed() const;
   bool isBottomClosed() const;
@@ -412,7 +434,12 @@ public:
   void dump() const;
 
 protected:
-  const LiveRange *getLiveRange(unsigned Reg) const;
+  void discoverLiveOut(unsigned Reg);
+  void discoverLiveIn(unsigned Reg);
+
+  /// \brief Get the SlotIndex for the first nondebug instruction including or
+  /// after the current position.
+  SlotIndex getCurrSlot() const;
 
   void increaseRegPressure(ArrayRef<unsigned> Regs);
   void decreaseRegPressure(ArrayRef<unsigned> Regs);