[NFC][regalloc] Separate iteration from AllocationOrder

This separates the two concerns - encapsulation of traversal order; and
iteration.

Differential Revision: https://reviews.llvm.org/D88256

lib/CodeGen/AllocationOrder.h

@@ -17,8 +17,8 @@
 #define LLVM_LIB_CODEGEN_ALLOCATIONORDER_H
 
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/MC/MCRegister.h"
 
 namespace llvm {
@@ -30,12 +30,52 @@ class LiveRegMatrix;
 class LLVM_LIBRARY_VISIBILITY AllocationOrder {
   const SmallVector<MCPhysReg, 16> Hints;
   ArrayRef<MCPhysReg> Order;
-  int Pos = 0;
-
-  // If HardHints is true, *only* Hints will be returned.
-  const bool HardHints;
+  // How far into the Order we can iterate. This is 0 if the AllocationOrder is
+  // constructed with HardHints = true, Order.size() otherwise. While
+  // technically a size_t, it will participate in comparisons with the
+  // Iterator's Pos, which must be signed, so it's typed here as signed, too, to
+  // avoid warnings and under the assumption that the size of Order is
+  // relatively small.
+  // IterationLimit defines an invalid iterator position.
+  const int IterationLimit;
 
 public:
+  /// Forward iterator for an AllocationOrder.
+  class Iterator final {
+    const AllocationOrder &AO;
+    int Pos = 0;
+
+  public:
+    Iterator(const AllocationOrder &AO, int Pos) : AO(AO), Pos(Pos) {}
+
+    /// Return true if the curent position is that of a preferred register.
+    bool isHint() const { return Pos < 0; }
+
+    /// Return the next physical register in the allocation order.
+    MCRegister operator*() const {
+      if (Pos < 0)
+        return AO.Hints.end()[Pos];
+      assert(Pos < AO.IterationLimit);
+      return AO.Order[Pos];
+    }
+
+    /// Advance the iterator to the next position. If that's past the Hints
+    /// list, advance to the first value that's not also in the Hints list.
+    Iterator &operator++() {
+      if (Pos < AO.IterationLimit)
+        ++Pos;
+      while (Pos >= 0 && Pos < AO.IterationLimit && AO.isHint(AO.Order[Pos]))
+        ++Pos;
+      return *this;
+    }
+
+    bool operator==(const Iterator &Other) const {
+      assert(&AO == &Other.AO);
+      return Pos == Other.Pos;
+    }
+
+    bool operator!=(const Iterator &Other) const { return !(*this == Other); }
+  };
 
   /// Create a new AllocationOrder for VirtReg.
   /// @param VirtReg      Virtual register to allocate for.
@@ -50,35 +90,26 @@ public:
   AllocationOrder(SmallVector<MCPhysReg, 16> &&Hints, ArrayRef<MCPhysReg> Order,
                   bool HardHints)
       : Hints(std::move(Hints)), Order(Order),
-        Pos(-static_cast<int>(this->Hints.size())), HardHints(HardHints) {}
+        IterationLimit(HardHints ? 0 : static_cast<int>(Order.size())) {}
+
+  Iterator begin() const {
+    return Iterator(*this, -(static_cast<int>(Hints.size())));
+  }
+
+  Iterator end() const { return Iterator(*this, IterationLimit); }
+
+  Iterator getOrderLimitEnd(unsigned OrderLimit) const {
+    assert(OrderLimit <= Order.size());
+    if (OrderLimit == 0)
+      return end();
+    Iterator Ret(*this,
+                 std::min(static_cast<int>(OrderLimit) - 1, IterationLimit));
+    return ++Ret;
+  }
 
   /// Get the allocation order without reordered hints.
   ArrayRef<MCPhysReg> getOrder() const { return Order; }
 
-  /// Return the next physical register in the allocation order, or 0.
-  /// It is safe to call next() again after it returned 0, it will keep
-  /// returning 0 until rewind() is called.
-  MCPhysReg next(unsigned Limit = 0) {
-    if (Pos < 0)
-      return Hints.end()[Pos++];
-    if (HardHints)
-      return 0;
-    if (!Limit)
-      Limit = Order.size();
-    while (Pos < int(Limit)) {
-      unsigned Reg = Order[Pos++];
-      if (!isHint(Reg))
-        return Reg;
-    }
-    return 0;
-  }
-
-  /// Start over from the beginning.
-  void rewind() { Pos = -int(Hints.size()); }
-
-  /// Return true if the last register returned from next() was a preferred register.
-  bool isHint() const { return Pos <= 0; }
-
   /// Return true if PhysReg is a preferred register.
   bool isHint(unsigned PhysReg) const { return is_contained(Hints, PhysReg); }
 };

lib/CodeGen/RegAllocBasic.cpp

@@ -261,7 +261,8 @@ Register RABasic::selectOrSplit(LiveInterval &VirtReg,
   // Check for an available register in this class.
   auto Order =
       AllocationOrder::create(VirtReg.reg(), *VRM, RegClassInfo, Matrix);
-  while (Register PhysReg = Order.next()) {
+  for (MCRegister PhysReg : Order) {
+    assert(PhysReg.isValid());
     // Check for interference in PhysReg
     switch (Matrix->checkInterference(VirtReg, PhysReg)) {
     case LiveRegMatrix::IK_Free:

lib/CodeGen/RegAllocGreedy.cpp

@@ -757,12 +757,17 @@ Register RAGreedy::tryAssign(LiveInterval &VirtReg,
                              AllocationOrder &Order,
                              SmallVectorImpl<Register> &NewVRegs,
                              const SmallVirtRegSet &FixedRegisters) {
-  Order.rewind();
   Register PhysReg;
-  while ((PhysReg = Order.next()))
-    if (!Matrix->checkInterference(VirtReg, PhysReg))
-      break;
-  if (!PhysReg || Order.isHint())
+  for (auto I = Order.begin(), E = Order.end(); I != E && !PhysReg; ++I) {
+    assert(*I);
+    if (!Matrix->checkInterference(VirtReg, *I)) {
+      if (I.isHint())
+        return *I;
+      else
+        PhysReg = *I;
+    }
+  }
+  if (!PhysReg.isValid())
     return PhysReg;
 
   // PhysReg is available, but there may be a better choice.
@@ -803,12 +808,12 @@ Register RAGreedy::tryAssign(LiveInterval &VirtReg,
 Register RAGreedy::canReassign(LiveInterval &VirtReg, Register PrevReg) {
   auto Order =
       AllocationOrder::create(VirtReg.reg(), *VRM, RegClassInfo, Matrix);
-  Register PhysReg;
-  while ((PhysReg = Order.next())) {
-    if (PhysReg == PrevReg)
+  MCRegister PhysReg;
+  for (auto I = Order.begin(), E = Order.end(); I != E && !PhysReg; ++I) {
+    if ((*I).id() == PrevReg.id())
       continue;
 
-    MCRegUnitIterator Units(PhysReg, TRI);
+    MCRegUnitIterator Units(*I, TRI);
     for (; Units.isValid(); ++Units) {
       // Instantiate a "subquery", not to be confused with the Queries array.
       LiveIntervalUnion::Query subQ(VirtReg, Matrix->getLiveUnions()[*Units]);
@@ -817,7 +822,7 @@ Register RAGreedy::canReassign(LiveInterval &VirtReg, Register PrevReg) {
     }
     // If no units have interference, break out with the current PhysReg.
     if (!Units.isValid())
-      break;
+      PhysReg = *I;
   }
   if (PhysReg)
     LLVM_DEBUG(dbgs() << "can reassign: " << VirtReg << " from "
@@ -1134,8 +1139,10 @@ unsigned RAGreedy::tryEvict(LiveInterval &VirtReg,
     }
   }
 
-  Order.rewind();
-  while (MCRegister PhysReg = Order.next(OrderLimit)) {
+  for (auto I = Order.begin(), E = Order.getOrderLimitEnd(OrderLimit); I != E;
+       ++I) {
+    MCRegister PhysReg = *I;
+    assert(PhysReg);
     if (TRI->getCostPerUse(PhysReg) >= CostPerUseLimit)
       continue;
     // The first use of a callee-saved register in a function has cost 1.
@@ -1156,7 +1163,7 @@ unsigned RAGreedy::tryEvict(LiveInterval &VirtReg,
     BestPhys = PhysReg;
 
     // Stop if the hint can be used.
-    if (Order.isHint())
+    if (I.isHint())
       break;
   }
 
@@ -1849,8 +1856,8 @@ unsigned RAGreedy::calculateRegionSplitCost(LiveInterval &VirtReg,
                                             unsigned &NumCands, bool IgnoreCSR,
                                             bool *CanCauseEvictionChain) {
   unsigned BestCand = NoCand;
-  Order.rewind();
-  while (unsigned PhysReg = Order.next()) {
+  for (MCPhysReg PhysReg : Order) {
+    assert(PhysReg);
     if (IgnoreCSR && isUnusedCalleeSavedReg(PhysReg))
       continue;
 
@@ -2288,8 +2295,8 @@ unsigned RAGreedy::tryLocalSplit(LiveInterval &VirtReg, AllocationOrder &Order,
     (1.0f / MBFI->getEntryFreq());
   SmallVector<float, 8> GapWeight;
 
-  Order.rewind();
-  while (unsigned PhysReg = Order.next()) {
+  for (MCPhysReg PhysReg : Order) {
+    assert(PhysReg);
     // Keep track of the largest spill weight that would need to be evicted in
     // order to make use of PhysReg between UseSlots[I] and UseSlots[I + 1].
     calcGapWeights(PhysReg, GapWeight);
@@ -2606,8 +2613,8 @@ unsigned RAGreedy::tryLastChanceRecoloring(LiveInterval &VirtReg,
   FixedRegisters.insert(VirtReg.reg());
   SmallVector<Register, 4> CurrentNewVRegs;
 
-  Order.rewind();
-  while (Register PhysReg = Order.next()) {
+  for (MCRegister PhysReg : Order) {
+    assert(PhysReg.isValid());
     LLVM_DEBUG(dbgs() << "Try to assign: " << VirtReg << " to "
                       << printReg(PhysReg, TRI) << '\n');
     RecoloringCandidates.clear();

unittests/CodeGen/AllocationOrderTest.cpp

@@ -12,11 +12,14 @@
 using namespace llvm;
 
 namespace {
-std::vector<MCPhysReg> loadOrder(AllocationOrder &O, unsigned Limit = 0) {
+std::vector<MCPhysReg> loadOrder(const AllocationOrder &O, unsigned Limit = 0) {
   std::vector<MCPhysReg> Ret;
-  O.rewind();
-  while (auto R = O.next(Limit))
-    Ret.push_back(R);
+  if (Limit == 0)
+    for (auto R : O)
+      Ret.push_back(R);
+  else
+    for (auto I = O.begin(), E = O.getOrderLimitEnd(Limit); I != E; ++I)
+      Ret.push_back(*I);
   return Ret;
 }
 } // namespace
@@ -48,6 +51,7 @@ TEST(AllocationOrderTest, LimitsBasic) {
   AllocationOrder O(std::move(Hints), Order, false);
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3, 4, 5, 6, 7}), loadOrder(O, 0));
   EXPECT_EQ((std::vector<MCPhysReg>{1, 2, 3, 4}), loadOrder(O, 1));
+  EXPECT_EQ(O.end(), O.getOrderLimitEnd(0));
 }
 
 TEST(AllocationOrderTest, LimitsDuplicates) {
@@ -96,19 +100,19 @@ TEST(AllocationOrderTest, IsHintTest) {
   SmallVector<MCPhysReg, 16> Hints = {1, 2, 3};
   SmallVector<MCPhysReg, 16> Order = {4, 1, 5, 6};
   AllocationOrder O(std::move(Hints), Order, false);
-  O.rewind();
-  auto V = O.next();
-  EXPECT_TRUE(O.isHint());
+  auto I = O.begin();
+  auto V = *I;
+  EXPECT_TRUE(I.isHint());
   EXPECT_EQ(V, 1U);
-  O.next();
-  EXPECT_TRUE(O.isHint());
-  O.next();
-  EXPECT_TRUE(O.isHint());
-  V = O.next();
-  EXPECT_FALSE(O.isHint());
+  ++I;
+  EXPECT_TRUE(I.isHint());
+  ++I;
+  EXPECT_TRUE(I.isHint());
+  V = *(++I);
+  EXPECT_FALSE(I.isHint());
   EXPECT_EQ(V, 4U);
-  V = O.next();
+  V = *(++I);
   EXPECT_TRUE(O.isHint(1));
-  EXPECT_FALSE(O.isHint());
+  EXPECT_FALSE(I.isHint());
   EXPECT_EQ(V, 5U);
 }