[CodeGen] Enhance `MachineInstrSpan` to allow the end of MBB to be used.

Summary:
- Explicitly specify the parent MBB to allow the end iterator to be
  used.

Reviewers: aprantl, MatzeB, craig.topper, qcolombet

Subscribers: arsenm, jvesely, nhaehnle, hiraditya, llvm-commits

Tags: #llvm

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

llvm-svn: 365240

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -910,11 +910,11 @@ class MachineInstrSpan {
   MachineBasicBlock::iterator I, B, E;
 
 public:
-  MachineInstrSpan(MachineBasicBlock::iterator I)
-    : MBB(*I->getParent()),
-      I(I),
-      B(I == MBB.begin() ? MBB.end() : std::prev(I)),
-      E(std::next(I)) {}
+  MachineInstrSpan(MachineBasicBlock::iterator I, MachineBasicBlock *BB)
+      : MBB(*BB), I(I), B(I == MBB.begin() ? MBB.end() : std::prev(I)),
+        E(std::next(I)) {
+    assert(I == BB->end() || I->getParent() == BB);
+  }
 
   MachineBasicBlock::iterator begin() {
     return B == MBB.end() ? MBB.begin() : std::next(B);

llvm/lib/CodeGen/InlineSpiller.cpp

@@ -833,7 +833,7 @@ foldMemoryOperand(ArrayRef<std::pair<MachineInstr *, unsigned>> Ops,
   if (FoldOps.empty())
     return false;
 
-  MachineInstrSpan MIS(MI);
+  MachineInstrSpan MIS(MI, MI->getParent());
 
   MachineInstr *FoldMI =
       LoadMI ? TII.foldMemoryOperand(*MI, FoldOps, *LoadMI, &LIS)
@@ -907,7 +907,7 @@ void InlineSpiller::insertReload(unsigned NewVReg,
                                  MachineBasicBlock::iterator MI) {
   MachineBasicBlock &MBB = *MI->getParent();
 
-  MachineInstrSpan MIS(MI);
+  MachineInstrSpan MIS(MI, &MBB);
   TII.loadRegFromStackSlot(MBB, MI, NewVReg, StackSlot,
                            MRI.getRegClass(NewVReg), &TRI);
 
@@ -937,7 +937,7 @@ void InlineSpiller::insertSpill(unsigned NewVReg, bool isKill,
                                  MachineBasicBlock::iterator MI) {
   MachineBasicBlock &MBB = *MI->getParent();
 
-  MachineInstrSpan MIS(MI);
+  MachineInstrSpan MIS(MI, &MBB);
   bool IsRealSpill = true;
   if (isFullUndefDef(*MI)) {
     // Don't spill undef value.

llvm/lib/Target/AMDGPU/SILowerSGPRSpills.cpp

@@ -92,7 +92,7 @@ static void insertCSRSaves(MachineBasicBlock &SaveBlock,
       // Insert the spill to the stack frame.
       unsigned Reg = CS.getReg();
 
-      MachineInstrSpan MIS(I);
+      MachineInstrSpan MIS(I, &SaveBlock);
       const TargetRegisterClass *RC = TRI->getMinimalPhysRegClass(Reg);
 
       TII.storeRegToStackSlot(SaveBlock, I, Reg, true, CS.getFrameIdx(), RC,

llvm/unittests/CodeGen/MachineInstrTest.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
@@ -273,6 +274,38 @@ TEST(MachineInstrPrintingTest, DebugLocPrinting) {
       StringRef(OS.str()).endswith("filename:1:5"));
 }
 
+TEST(MachineInstrSpan, DistanceBegin) {
+  auto MF = createMachineFunction();
+  auto MBB = MF->CreateMachineBasicBlock();
+
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MII = MBB->begin();
+  MachineInstrSpan MIS(MII, MBB);
+  ASSERT_TRUE(MIS.empty());
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MII, MI);
+  ASSERT_TRUE(std::distance(MIS.begin(), MII) == 1);
+}
+
+TEST(MachineInstrSpan, DistanceEnd) {
+  auto MF = createMachineFunction();
+  auto MBB = MF->CreateMachineBasicBlock();
+
+  MCInstrDesc MCID = {0, 0,       0,       0,       0, 0,
+                      0, nullptr, nullptr, nullptr, 0, nullptr};
+
+  auto MII = MBB->end();
+  MachineInstrSpan MIS(MII, MBB);
+  ASSERT_TRUE(MIS.empty());
+
+  auto MI = MF->CreateMachineInstr(MCID, DebugLoc());
+  MBB->insert(MII, MI);
+  ASSERT_TRUE(std::distance(MIS.begin(), MII) == 1);
+}
+
 static_assert(is_trivially_copyable<MCOperand>::value, "trivially copyable");
 
 } // end namespace