3 files changed, 56 insertions, 8 deletions

diff --git a/compiler/optimizing/register_allocator.cc b/compiler/optimizing/register_allocator.cc
index f53f846326..925099ade6 100644
--- a/compiler/optimizing/register_allocator.cc
+++ b/compiler/optimizing/register_allocator.cc
@@ -775,7 +775,7 @@ bool RegisterAllocator::TryAllocateFreeReg(LiveInterval* current) {
     } else if (current->IsLowInterval()) {
       reg = FindAvailableRegisterPair(free_until, current->GetStart());
     } else {
-      reg = FindAvailableRegister(free_until);
+      reg = FindAvailableRegister(free_until, current);
     }
   }
 
@@ -839,14 +839,52 @@ int RegisterAllocator::FindAvailableRegisterPair(size_t* next_use, size_t starti
   return reg;
 }
 
-int RegisterAllocator::FindAvailableRegister(size_t* next_use) const {
+bool RegisterAllocator::IsCallerSaveRegister(int reg) const {
+  return processing_core_registers_
+      ? !codegen_->IsCoreCalleeSaveRegister(reg)
+      : !codegen_->IsFloatingPointCalleeSaveRegister(reg);
+}
+
+int RegisterAllocator::FindAvailableRegister(size_t* next_use, LiveInterval* current) const {
+  // We special case intervals that do not span a safepoint to try to find a caller-save
+  // register if one is available. We iterate from 0 to the number of registers,
+  // so if there are caller-save registers available at the end, we continue the iteration.
+  bool prefers_caller_save = !current->HasWillCallSafepoint();
   int reg = kNoRegister;
-  // Pick the register that is used the last.
   for (size_t i = 0; i < number_of_registers_; ++i) {
-    if (IsBlocked(i)) continue;
-    if (reg == kNoRegister || next_use[i] > next_use[reg]) {
+    if (IsBlocked(i)) {
+      // Register cannot be used. Continue.
+      continue;
+    }
+
+    // Best case: we found a register fully available.
+    if (next_use[i] == kMaxLifetimePosition) {
+      if (prefers_caller_save && !IsCallerSaveRegister(i)) {
+        // We can get shorter encodings on some platforms by using
+        // small register numbers. So only update the candidate if the previous
+        // one was not available for the whole method.
+        if (reg == kNoRegister || next_use[reg] != kMaxLifetimePosition) {
+          reg = i;
+        }
+        // Continue the iteration in the hope of finding a caller save register.
+        continue;
+      } else {
+        reg = i;
+        // We know the register is good enough. Return it.
+        break;
+      }
+    }
+
+    // If we had no register before, take this one as a reference.
+    if (reg == kNoRegister) {
       reg = i;
-      if (next_use[i] == kMaxLifetimePosition) break;
+      continue;
+    }
+
+    // Pick the register that is used the last.
+    if (next_use[i] > next_use[reg]) {
+      reg = i;
+      continue;
     }
   }
   return reg;
@@ -971,7 +1009,7 @@ bool RegisterAllocator::AllocateBlockedReg(LiveInterval* current) {
       || (first_use >= next_use[GetHighForLowRegister(reg)]);
   } else {
     DCHECK(!current->IsHighInterval());
-    reg = FindAvailableRegister(next_use);
+    reg = FindAvailableRegister(next_use, current);
     should_spill = (first_use >= next_use[reg]);
   }
 
diff --git a/compiler/optimizing/register_allocator.h b/compiler/optimizing/register_allocator.h
index dc9c708eea..6d5bfc3f0d 100644
--- a/compiler/optimizing/register_allocator.h
+++ b/compiler/optimizing/register_allocator.h
@@ -140,7 +140,8 @@ class RegisterAllocator {
   void DumpInterval(std::ostream& stream, LiveInterval* interval) const;
   void DumpAllIntervals(std::ostream& stream) const;
   int FindAvailableRegisterPair(size_t* next_use, size_t starting_at) const;
-  int FindAvailableRegister(size_t* next_use) const;
+  int FindAvailableRegister(size_t* next_use, LiveInterval* current) const;
+  bool IsCallerSaveRegister(int reg) const;
 
   // Try splitting an active non-pair or unaligned pair interval at the given `position`.
   // Returns whether it was successful at finding such an interval.
diff --git a/compiler/optimizing/ssa_liveness_analysis.h b/compiler/optimizing/ssa_liveness_analysis.h
index bd55e9fc7f..ce4bbd4fc7 100644
--- a/compiler/optimizing/ssa_liveness_analysis.h
+++ b/compiler/optimizing/ssa_liveness_analysis.h
@@ -542,6 +542,15 @@ class LiveInterval : public ArenaObject<kArenaAllocMisc> {
     return defined_by_;
   }
 
+  bool HasWillCallSafepoint() const {
+    for (SafepointPosition* safepoint = first_safepoint_;
+         safepoint != nullptr;
+         safepoint = safepoint->GetNext()) {
+      if (safepoint->GetLocations()->WillCall()) return true;
+    }
+    return false;
+  }
+
   SafepointPosition* FindSafepointJustBefore(size_t position) const {
     for (SafepointPosition* safepoint = first_safepoint_, *previous = nullptr;
          safepoint != nullptr;