Index: lib/CodeGen/MachineCopyPropagation.cpp
===================================================================
--- lib/CodeGen/MachineCopyPropagation.cpp
+++ lib/CodeGen/MachineCopyPropagation.cpp
@@ -75,40 +75,36 @@
 namespace {
 
 class CopyTracker {
-  using RegList = SmallVector<unsigned, 4>;
-  using SourceMap = DenseMap<unsigned, RegList>;
-  using Reg2MIMap = DenseMap<unsigned, MachineInstr *>;
+  struct CopyInfo {
+    MachineInstr *MI;
+    SmallVector<unsigned, 4> DefRegs;
+    bool Avail;
+  };
 
-  /// Def -> available copies map.
-  Reg2MIMap AvailCopyMap;
-
-  /// Def -> copies map.
-  Reg2MIMap CopyMap;
-
-  /// Src -> Def map
-  SourceMap SrcMap;
+  DenseMap<unsigned, CopyInfo> Copies;
 
 public:
   /// Mark all of the given registers and their subregisters as unavailable for
   /// copying.
-  void markRegsUnavailable(const RegList &Regs, const TargetRegisterInfo &TRI) {
+  void markRegsUnavailable(ArrayRef<unsigned> Regs,
+                           const TargetRegisterInfo &TRI) {
     for (unsigned Reg : Regs) {
       // Source of copy is no longer available for propagation.
-      for (MCSubRegIterator SR(Reg, &TRI, true); SR.isValid(); ++SR)
-        AvailCopyMap.erase(*SR);
+      for (MCRegUnitIterator RUI(Reg, &TRI); RUI.isValid(); ++RUI) {
+        auto CI = Copies.find(*RUI);
+        if (CI != Copies.end())
+          CI->second.Avail = false;
+      }
     }
   }
 
   /// Clobber a single register, removing it from the tracker's copy maps.
   void clobberRegister(unsigned Reg, const TargetRegisterInfo &TRI) {
-    for (MCRegAliasIterator AI(Reg, &TRI, true); AI.isValid(); ++AI) {
-      CopyMap.erase(*AI);
-      AvailCopyMap.erase(*AI);
-
-      SourceMap::iterator SI = SrcMap.find(*AI);
-      if (SI != SrcMap.end()) {
-        markRegsUnavailable(SI->second, TRI);
-        SrcMap.erase(SI);
+    for (MCRegUnitIterator RUI(Reg, &TRI); RUI.isValid(); ++RUI) {
+      auto I = Copies.find(*RUI);
+      if (I != Copies.end()) {
+        markRegsUnavailable(I->second.DefRegs, TRI);
+        Copies.erase(I);
       }
     }
   }
@@ -121,39 +117,49 @@
     unsigned Src = Copy->getOperand(1).getReg();
 
     // Remember Def is defined by the copy.
-    for (MCSubRegIterator SR(Def, &TRI, /*IncludeSelf=*/true); SR.isValid();
-         ++SR) {
-      CopyMap[*SR] = Copy;
-      AvailCopyMap[*SR] = Copy;
-    }
+    for (MCRegUnitIterator RUI(Def, &TRI); RUI.isValid(); ++RUI)
+      Copies[*RUI] = {Copy, {}, true};
 
     // Remember source that's copied to Def. Once it's clobbered, then
     // it's no longer available for copy propagation.
-    RegList &DestList = SrcMap[Src];
-    if (!is_contained(DestList, Def))
-      DestList.push_back(Def);
+    for (MCRegUnitIterator RUI(Src, &TRI); RUI.isValid(); ++RUI) {
+      auto I = Copies.insert({*RUI, {nullptr, {}, false}});
+      auto &Copy = I.first->second;
+      if (!is_contained(Copy.DefRegs, Def))
+        Copy.DefRegs.push_back(Def);
+    }
+  }
+
+  bool hasAvailableCopies() {
+    // TODO: This is conservatively correct, but is it useful?
+    return !Copies.empty();
   }
 
-  bool hasAvailableCopies() { return !AvailCopyMap.empty(); }
+  MachineInstr *findCopyForUnit(unsigned RegUnit, const TargetRegisterInfo &TRI,
+                         bool MustBeAvailable = false) {
+    auto CI = Copies.find(RegUnit);
+    if (CI == Copies.end())
+      return nullptr;
+    if (MustBeAvailable && !CI->second.Avail)
+      return nullptr;
+    return CI->second.MI;
+  }
 
-  MachineInstr *findAvailCopy(unsigned Reg) {
-    auto CI = AvailCopyMap.find(Reg);
-    if (CI != AvailCopyMap.end())
-      return CI->second;
-    return nullptr;
+  MachineInstr *findCopy(unsigned Reg, const TargetRegisterInfo &TRI,
+                         bool MustBeAvailable = false) {
+    MCRegUnitIterator RUI(Reg, &TRI);
+    MachineInstr *MI = findCopyForUnit(*RUI, TRI, MustBeAvailable);
+    if (!MI || !TRI.isSubRegisterEq(MI->getOperand(0).getReg(), Reg))
+      return nullptr;
+    return MI;
   }
 
-  MachineInstr *findCopy(unsigned Reg) {
-    auto CI = CopyMap.find(Reg);
-    if (CI != CopyMap.end())
-      return CI->second;
-    return nullptr;
+  MachineInstr *findAvailCopy(unsigned Reg, const TargetRegisterInfo &TRI) {
+    return findCopy(Reg, TRI, /*MustBeAvailable=*/true);
   }
 
   void clear() {
-    AvailCopyMap.clear();
-    CopyMap.clear();
-    SrcMap.clear();
+    Copies.clear();
   }
 };
 
@@ -211,8 +217,8 @@
 void MachineCopyPropagation::ReadRegister(unsigned Reg) {
   // If 'Reg' is defined by a copy, the copy is no longer a candidate
   // for elimination.
-  for (MCRegAliasIterator AI(Reg, TRI, true); AI.isValid(); ++AI) {
-    if (MachineInstr *Copy = Tracker.findCopy(*AI)) {
+  for (MCRegUnitIterator RUI(Reg, TRI); RUI.isValid(); ++RUI) {
+    if (MachineInstr *Copy = Tracker.findCopyForUnit(*RUI, *TRI)) {
       LLVM_DEBUG(dbgs() << "MCP: Copy is used - not dead: "; Copy->dump());
       MaybeDeadCopies.remove(Copy);
     }
@@ -250,7 +256,7 @@
     return false;
 
   // Search for an existing copy.
-  MachineInstr *PrevCopy = Tracker.findAvailCopy(Def);
+  MachineInstr *PrevCopy = Tracker.findAvailCopy(Def, *TRI);
   if (!PrevCopy)
     return false;
 
@@ -384,7 +390,7 @@
     if (!MOUse.isRenamable())
       continue;
 
-    MachineInstr *Copy = Tracker.findAvailCopy(MOUse.getReg());
+    MachineInstr *Copy = Tracker.findAvailCopy(MOUse.getReg(), *TRI);
     if (!Copy)
       continue;