2 files changed, 17 insertions, 8 deletions
diff --git a/llvm/include/llvm/CodeGen/MachineInstrBundleIterator.h b/llvm/include/llvm/CodeGen/MachineInstrBundleIterator.h
index 8ee3a16b36a..f493dbc5e36 100644
--- a/llvm/include/llvm/CodeGen/MachineInstrBundleIterator.h
+++ b/llvm/include/llvm/CodeGen/MachineInstrBundleIterator.h
@@ -37,6 +37,13 @@ public:
   typedef typename instr_iterator::const_pointer const_pointer;
   typedef typename instr_iterator::const_reference const_reference;
 
+private:
+  typedef typename std::remove_const<value_type>::type nonconst_value_type;
+  typedef ilist_node<nonconst_value_type> node_type;
+  typedef ilist_iterator<nonconst_value_type> nonconst_instr_iterator;
+  typedef MachineInstrBundleIterator<nonconst_value_type> nonconst_iterator;
+
+public:
   MachineInstrBundleIterator(instr_iterator MI) : MII(MI) {}
 
   MachineInstrBundleIterator(reference MI) : MII(MI) {
@@ -130,6 +137,12 @@ public:
   }
 
   instr_iterator getInstrIterator() const { return MII; }
+
+  nonconst_iterator getNonConstIterator() const {
+    if (auto *N = const_cast<node_type *>(MII.getNodePtr()))
+      return nonconst_iterator(nonconst_instr_iterator(*N));
+    return nonconst_iterator();
+  }
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/CodeGen/MachineScheduler.cpp b/llvm/lib/CodeGen/MachineScheduler.cpp
index 01ef9d834b6..e91fa051f36 100644
--- a/llvm/lib/CodeGen/MachineScheduler.cpp
+++ b/llvm/lib/CodeGen/MachineScheduler.cpp
@@ -251,8 +251,8 @@ priorNonDebug(MachineBasicBlock::const_iterator I,
 static MachineBasicBlock::iterator
 priorNonDebug(MachineBasicBlock::iterator I,
               MachineBasicBlock::const_iterator Beg) {
-  return const_cast<MachineInstr*>(
-    &*priorNonDebug(MachineBasicBlock::const_iterator(I), Beg));
+  return priorNonDebug(MachineBasicBlock::const_iterator(I), Beg)
+      .getNonConstIterator();
 }
 
 /// If this iterator is a debug value, increment until reaching the End or a
@@ -271,12 +271,8 @@ nextIfDebug(MachineBasicBlock::const_iterator I,
 static MachineBasicBlock::iterator
 nextIfDebug(MachineBasicBlock::iterator I,
             MachineBasicBlock::const_iterator End) {
-  // Cast the return value to nonconst MachineInstr, then cast to an
-  // instr_iterator, which does not check for null, finally return a
-  // bundle_iterator.
-  return MachineBasicBlock::instr_iterator(
-    const_cast<MachineInstr*>(
-      &*nextIfDebug(MachineBasicBlock::const_iterator(I), End)));
+  return nextIfDebug(MachineBasicBlock::const_iterator(I), End)
+      .getNonConstIterator();
 }
 
 /// Instantiate a ScheduleDAGInstrs that will be owned by the caller.