From 0aef948a1e917ba7560ed2279c493ab0468b5fe5 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Thu, 6 Mar 2025 17:24:59 -0800
Subject: [PATCH] Load an immediate into a register in the backend
 (https://github.com/Shopify/zjit/pull/38)

---
 zjit/src/backend/x86_64/mod.rs | 16 +++++++++++++---
 zjit/src/codegen.rs            | 16 ++--------------
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
index b40bfbf797..078459540c 100644
--- a/zjit/src/backend/x86_64/mod.rs
+++ b/zjit/src/backend/x86_64/mod.rs
@@ -217,6 +217,10 @@ impl Assembler
                                 (Opnd::Mem(_) | Opnd::Reg(_), _) => {
                                     *left = asm.load(*left);
                                 },
+                                // The first operand can't be an immediate value
+                                (Opnd::Value(_), _) => {
+                                    *left = asm.load(*left);
+                                }
                                 _ => {}
                             };
 
@@ -244,9 +248,15 @@ impl Assembler
                     }
                 },
                 Insn::Test { left, right } => {
-                    if let (Opnd::Mem(_), Opnd::Mem(_)) = (&left, &right) {
-                        let loaded = asm.load(*right);
-                        *right = loaded;
+                    match (&left, &right) {
+                        (Opnd::Mem(_), Opnd::Mem(_)) => {
+                            *right = asm.load(*right);
+                        }
+                        // The first operand can't be an immediate value
+                        (Opnd::UImm(_) | Opnd::Imm(_), _) => {
+                            *left = asm.load(*left);
+                        }
+                        _ => {}
                     }
                     asm.push_insn(insn);
                 },
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index f3a9b68e4b..ad97748cf9 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -146,14 +146,8 @@ fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, left: InsnId, right:
     let left_opnd = jit.get_opnd(left)?;
     let right_opnd = jit.get_opnd(right)?;
 
-    // Load left into a register if left is a constant. The backend doesn't support sub(imm, imm).
-    let left_reg = match left_opnd {
-        Opnd::Value(_) => asm.load(left_opnd),
-        _ => left_opnd,
-    };
-
     // Add arg0 + arg1 and test for overflow
-    let left_untag = asm.sub(left_reg, Opnd::Imm(1));
+    let left_untag = asm.sub(left_opnd, Opnd::Imm(1));
     let out_val = asm.add(left_untag, right_opnd);
     asm.jo(Target::SideExit(state.clone()));
 
@@ -164,14 +158,8 @@ fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, left: InsnId, right:
 fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: InsnId, guard_type: Type, state: &FrameState) -> Option<lir::Opnd> {
     let opnd = jit.get_opnd(val)?;
     if guard_type.is_subtype(Fixnum) {
-        // Load opnd into a register if opnd is a constant. The backend doesn't support test(imm, imm) yet.
-        let opnd_reg = match opnd {
-            Opnd::Value(_) => asm.load(opnd),
-            _ => opnd,
-        };
-
         // Check if opnd is Fixnum
-        asm.test(opnd_reg, Opnd::UImm(RUBY_FIXNUM_FLAG as u64));
+        asm.test(opnd, Opnd::UImm(RUBY_FIXNUM_FLAG as u64));
         asm.jz(Target::SideExit(state.clone()));
     } else {
         unimplemented!("unsupported type: {guard_type}");