1 files changed, 22 insertions, 8 deletions
diff --git a/lib/CodeGen/LiveRegUnits.cpp b/lib/CodeGen/LiveRegUnits.cpp
index f9ba4ffa6527..9f28db6287ba 100644
--- a/lib/CodeGen/LiveRegUnits.cpp
+++ b/lib/CodeGen/LiveRegUnits.cpp
@@ -19,8 +19,8 @@
 #include "llvm/CodeGen/MachineInstrBundle.h"
 #include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
 #include "llvm/MC/MCRegisterInfo.h"
-#include "llvm/Target/TargetRegisterInfo.h"
 
 using namespace llvm;
 
@@ -97,23 +97,37 @@ static void addCalleeSavedRegs(LiveRegUnits &LiveUnits,
     LiveUnits.addReg(*CSR);
 }
 
-/// Adds pristine registers to the given \p LiveUnits. Pristine registers are
-/// callee saved registers that are unused in the function.
-static void addPristines(LiveRegUnits &LiveUnits, const MachineFunction &MF) {
+void LiveRegUnits::addPristines(const MachineFunction &MF) {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
   if (!MFI.isCalleeSavedInfoValid())
     return;
+  /// This function will usually be called on an empty object, handle this
+  /// as a special case.
+  if (empty()) {
+    /// Add all callee saved regs, then remove the ones that are saved and
+    /// restored.
+    addCalleeSavedRegs(*this, MF);
+    /// Remove the ones that are not saved/restored; they are pristine.
+    for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
+      removeReg(Info.getReg());
+    return;
+  }
+  /// If a callee-saved register that is not pristine is already present
+  /// in the set, we should make sure that it stays in it. Precompute the
+  /// set of pristine registers in a separate object.
   /// Add all callee saved regs, then remove the ones that are saved+restored.
-  addCalleeSavedRegs(LiveUnits, MF);
+  LiveRegUnits Pristine(*TRI);
+  addCalleeSavedRegs(Pristine, MF);
   /// Remove the ones that are not saved/restored; they are pristine.
   for (const CalleeSavedInfo &Info : MFI.getCalleeSavedInfo())
-    LiveUnits.removeReg(Info.getReg());
+    Pristine.removeReg(Info.getReg());
+  addUnits(Pristine.getBitVector());
 }
 
 void LiveRegUnits::addLiveOuts(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
   if (!MBB.succ_empty()) {
-    addPristines(*this, MF);
+    addPristines(MF);
     // To get the live-outs we simply merge the live-ins of all successors.
     for (const MachineBasicBlock *Succ : MBB.successors())
       addBlockLiveIns(*this, *Succ);
@@ -127,6 +141,6 @@ void LiveRegUnits::addLiveOuts(const MachineBasicBlock &MBB) {
 
 void LiveRegUnits::addLiveIns(const MachineBasicBlock &MBB) {
   const MachineFunction &MF = *MBB.getParent();
-  addPristines(*this, MF);
+  addPristines(MF);
   addBlockLiveIns(*this, MBB);
 }