From 8d7861e3daf64e0bd30b2f9fe56f94eadfde5d3f Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Fri, 4 Aug 2023 14:57:56 -0400
Subject: [PATCH] YJIT: expand bitwise shift support in x86 assembler (#8174)

---
 yjit/src/asm/x86_64/mod.rs   | 40 ++++++++++++++++++++++--------------
 yjit/src/asm/x86_64/tests.rs |  1 +
 2 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index 0bb1a6381f..45628dc084 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/yjit/src/asm/x86_64/mod.rs
@@ -1223,8 +1223,8 @@ pub fn ret(cb: &mut CodeBlock) {
     cb.write_byte(0xC3);
 }
 
-// Encode a single-operand shift instruction
-fn write_shift(cb: &mut CodeBlock, op_mem_one_pref: u8, _op_mem_cl_pref: u8, op_mem_imm_pref: u8, op_ext: Option<u8>, opnd0: X86Opnd, opnd1: X86Opnd) {
+// Encode a bitwise shift instruction
+fn write_shift(cb: &mut CodeBlock, op_mem_one_pref: u8, op_mem_cl_pref: u8, op_mem_imm_pref: u8, op_ext: u8, opnd0: X86Opnd, opnd1: X86Opnd) {
     assert!(matches!(opnd0, X86Opnd::Reg(_) | X86Opnd::Mem(_)));
 
     // Check the size of opnd0
@@ -1234,16 +1234,26 @@ fn write_shift(cb: &mut CodeBlock, op_mem_one_pref: u8, _op_mem_cl_pref: u8, op_
     let sz_pref = opnd_size == 16;
     let rex_w = opnd_size == 64;
 
-    if let X86Opnd::UImm(imm) = opnd1 {
-        if imm.value == 1 {
-            write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, op_ext, &[op_mem_one_pref]);
-        } else {
-            assert!(imm.num_bits <= 8);
-            write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, op_ext, &[op_mem_imm_pref]);
-            cb.write_byte(imm.value as u8);
+    match opnd1 {
+        X86Opnd::UImm(imm) => {
+            if imm.value == 1 {
+                write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, Some(op_ext), &[op_mem_one_pref]);
+            } else {
+                assert!(imm.num_bits <= 8);
+                write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, Some(op_ext), &[op_mem_imm_pref]);
+                cb.write_byte(imm.value as u8);
+            }
+        }
+
+        X86Opnd::Reg(reg) => {
+            // We can only use CL/RCX as the shift amount
+            assert!(reg.reg_no == RCX_REG.reg_no);
+            write_rm(cb, sz_pref, rex_w, X86Opnd::None, opnd0, Some(op_ext), &[op_mem_cl_pref]);
+        }
+
+        _ => {
+            unreachable!();
         }
-    } else {
-        unreachable!();
     }
 }
 
@@ -1254,7 +1264,7 @@ pub fn sal(cb: &mut CodeBlock, opnd0: X86Opnd, opnd1: X86Opnd) {
         0xD1, // opMemOnePref,
         0xD3, // opMemClPref,
         0xC1, // opMemImmPref,
-        Some(0x04),
+        0x04,
         opnd0,
         opnd1
     );
@@ -1267,7 +1277,7 @@ pub fn sar(cb: &mut CodeBlock, opnd0: X86Opnd, opnd1: X86Opnd) {
         0xD1, // opMemOnePref,
         0xD3, // opMemClPref,
         0xC1, // opMemImmPref,
-        Some(0x07),
+        0x07,
         opnd0,
         opnd1
     );
@@ -1280,7 +1290,7 @@ pub fn shl(cb: &mut CodeBlock, opnd0: X86Opnd, opnd1: X86Opnd) {
         0xD1, // opMemOnePref,
         0xD3, // opMemClPref,
         0xC1, // opMemImmPref,
-        Some(0x04),
+        0x04,
         opnd0,
         opnd1
     );
@@ -1293,7 +1303,7 @@ pub fn shr(cb: &mut CodeBlock, opnd0: X86Opnd, opnd1: X86Opnd) {
         0xD1, // opMemOnePref,
         0xD3, // opMemClPref,
         0xC1, // opMemImmPref,
-        Some(0x05),
+        0x05,
         opnd0,
         opnd1
     );
diff --git a/yjit/src/asm/x86_64/tests.rs b/yjit/src/asm/x86_64/tests.rs
index 2b1775452e..62f8a86230 100644
--- a/yjit/src/asm/x86_64/tests.rs
+++ b/yjit/src/asm/x86_64/tests.rs
@@ -340,6 +340,7 @@ fn test_sal() {
     check_bytes("d1e1", |cb| sal(cb, ECX, uimm_opnd(1)));
     check_bytes("c1e505", |cb| sal(cb, EBP, uimm_opnd(5)));
     check_bytes("d1642444", |cb| sal(cb, mem_opnd(32, RSP, 68), uimm_opnd(1)));
+    check_bytes("48d3e1", |cb| sal(cb, RCX, CL));
 }
 
 #[test]