From d4a154f366e9163d05b66697ea2b1ed4095b22a9 Mon Sep 17 00:00:00 2001
From: Laurence Tratt <laurie@tratt.net>
Date: Fri, 7 Feb 2025 15:27:02 +0000
Subject: [PATCH] Don't overwrite registers that we may still need to move.

---
 .../compile/jitc_yk/codegen/x64/lsregalloc.rs | 41 ++++++++++++-------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs b/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs
index 6883e0768..25ee967d7 100644
--- a/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs
+++ b/ykrt/src/compile/jitc_yk/codegen/x64/lsregalloc.rs
@@ -628,23 +628,26 @@ impl LSRegAlloc<'_> {
                                         &GPConstraint::Input { op: rhs_op, .. }
                                         | &GPConstraint::InputOutput { op: rhs_op, .. },
                                     ) => {
-                                        if lhs_op == rhs_op {
-                                            changed = true;
-                                            self.copy_gp_reg(asm, old_reg, reg);
-                                        } else {
-                                            if let RegState::Empty =
-                                                self.gp_reg_states[usize::from(reg.code())]
-                                            {
+                                        if let RegState::Empty =
+                                            self.gp_reg_states[usize::from(reg.code())]
+                                        {
+                                            if lhs_op == rhs_op {
+                                                self.copy_gp_reg(asm, old_reg, reg);
+                                            } else {
                                                 self.move_gp_reg(asm, old_reg, reg);
-                                            } else if let RegState::Empty =
-                                                self.gp_reg_states[usize::from(old_reg.code())]
-                                            {
-                                                self.move_gp_reg(asm, reg, old_reg);
+                                            }
+                                        } else if let RegState::Empty =
+                                            self.gp_reg_states[usize::from(old_reg.code())]
+                                        {
+                                            if lhs_op == rhs_op {
+                                                self.copy_gp_reg(asm, reg, old_reg);
                                             } else {
-                                                self.swap_gp_reg(asm, reg, old_reg);
+                                                self.move_gp_reg(asm, reg, old_reg);
                                             }
-                                            changed = true;
+                                        } else {
+                                            self.swap_gp_reg(asm, reg, old_reg);
                                         }
+                                        changed = true;
                                     }
                                     (
                                         &GPConstraint::Input { .. }
@@ -653,8 +656,18 @@ impl LSRegAlloc<'_> {
                                         | GPConstraint::Clobber { .. }
                                         | GPConstraint::Temporary,
                                     ) => {
+                                        if let RegState::Empty =
+                                            self.gp_reg_states[usize::from(reg.code())]
+                                        {
+                                            self.move_gp_reg(asm, old_reg, reg);
+                                        } else if let RegState::Empty =
+                                            self.gp_reg_states[usize::from(old_reg.code())]
+                                        {
+                                            self.move_gp_reg(asm, reg, old_reg);
+                                        } else {
+                                            self.swap_gp_reg(asm, reg, old_reg);
+                                        }
                                         changed = true;
-                                        self.move_gp_reg(asm, old_reg, reg);
                                     }
                                     _ => (),
                                 }