Fix untracked registers in stackmaps.

This commit fixes yet another bug in stackmaps that leads to registers being
untracked. The bug went unnoticed until the recent entryblock change exposed it.
The fix consists of two changes, allowing spill mappings to be applied
transitively and replacing killed registers in the stackmap with non-killed
registers holding the same value. See comments for more detail.

Author: ptersilie

llvm/lib/CodeGen/StackMaps.cpp

@@ -305,8 +305,22 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
     // each one of them. Note, that the stackmap may track either of %rbx or
     // %rcx, resulting in different ways below to retrieve the mappings.
     int ExtraReg = 0;
+    Register R = MOI->getReg();
     if (MOI->isReg()) {
-      Register R = MOI->getReg();
+      if (MOI->isKill()) {
+        // There's no point in tracking a killed register. Instead, try and
+        // find a copy of the register and use that to increase the likelyhood
+        // of us tracking that value.
+        // YKFIXME: This is likely only a temporary fix. In the future we might
+        // want to allow stackmaps to track arbitrarily many locations per live
+        // variable. Currently, we can only track two.
+        for (auto I : SpillOffsets) {
+          if (I.second == R) {
+            R = I.first;
+            break;
+          }
+        }
+      }
       if (SpillOffsets.count(R) > 0) {
         RHS = SpillOffsets[R];
         assert(SpillOffsets[R] != 0);
@@ -333,9 +347,9 @@ StackMaps::parseOperand(MachineInstr::const_mop_iterator MOI,
       }
     }
 
-    unsigned DwarfRegNum = getDwarfRegNum(MOI->getReg(), TRI);
+    unsigned DwarfRegNum = getDwarfRegNum(R, TRI);
     unsigned LLVMRegNum = *TRI->getLLVMRegNum(DwarfRegNum, false);
-    unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, MOI->getReg());
+    unsigned SubRegIdx = TRI->getSubRegIndex(LLVMRegNum, R);
     if (SubRegIdx)
       Offset = TRI->getSubRegIdxOffset(SubRegIdx);
 

llvm/lib/Target/X86/X86AsmPrinter.cpp

@@ -65,7 +65,21 @@ void clearRhs(Register Reg, std::map<Register, int64_t> &SpillMap) {
   auto I = SpillMap.begin();
   while (I != SpillMap.end()) {
     if (I->second == Reg) {
-      I = SpillMap.erase(I);
+      // If there's a mapping A => B, where B has been reassigned and has a
+      // mapping B => C, then transitively apply the mapping so that A => C.
+      // This makes sure we don't remove mappings (especially to stack slots)
+      // when a register is reassigned. Example:
+      //   store $r13 in [$rbp - 8]
+      //   $rcx = $r13
+      //   $r13 = $rbx
+      // Upon assigning `$r13 = $rbx`, we apply `$r13`'s previous mapping to
+      // `$rcx` so that `SpillMap[$rcx] = -8`.
+      if(SpillMap.count(Reg) > 0) {
+        SpillMap[I->first] = SpillMap[Reg];
+        ++I;
+      } else {
+        I = SpillMap.erase(I);
+      }
     } else {
       ++I;
     }
@@ -99,13 +113,14 @@ void processInstructions(
       const MachineOperand Rhs = Instr.getOperand(1);
       assert(Lhs.isReg() && "Is register.");
       assert(Rhs.isReg() && "Is register.");
+      // Reassigning a new value to LHS means any mappings to Lhs are now void
+      // and need to be removed. We need to do this before updating the
+      // mapping, so the transitive property of the SpillMap isn't violated
+      // (see `clearRhs` for more info).
+      clearRhs(Lhs.getReg(), SpillMap);
       SpillMap[Lhs.getReg()] = Rhs.getReg();
       // YKFIXME: If the `mov` instruction has a killed-flag, remove the
       // register from the map.
-
-      // Reassigning a new value to Lhs means any mappings to Lhs are now void
-      // and need to be removed.
-      clearRhs(Lhs.getReg(), SpillMap);
       continue;
     }