YJIT: implement `expandarray_rhs_too_small` case (#8153)

* YJIT: handle expandarray_rhs_too_small case

* YJIT: fix csel bug in x86 backend, add test

* Remove commented out lines

2 files changed, 59 insertions, 18 deletions

diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 0cc276fca1..1ae5ee7477 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -428,11 +428,31 @@ impl Assembler
             }
         }
 
-        fn emit_csel(cb: &mut CodeBlock, truthy: Opnd, falsy: Opnd, out: Opnd, cmov_fn: fn(&mut CodeBlock, X86Opnd, X86Opnd)) {
-            if out != truthy {
-                mov(cb, out.into(), truthy.into());
+        fn emit_csel(
+            cb: &mut CodeBlock,
+            truthy: Opnd,
+            falsy: Opnd,
+            out: Opnd,
+            cmov_fn: fn(&mut CodeBlock, X86Opnd, X86Opnd),
+            cmov_neg: fn(&mut CodeBlock, X86Opnd, X86Opnd)){
+
+            // Assert that output is a register
+            out.unwrap_reg();
+
+            // If the truthy value is a memory operand
+            if let Opnd::Mem(_) = truthy {
+                if out != falsy {
+                    mov(cb, out.into(), falsy.into());
+                }
+
+                cmov_fn(cb, out.into(), truthy.into());
+            } else {
+                if out != truthy {
+                    mov(cb, out.into(), truthy.into());
+                }
+
+                cmov_neg(cb, out.into(), falsy.into());
             }
-            cmov_fn(cb, out.into(), falsy.into());
         }
 
         //dbg!(&self.insns);
@@ -724,28 +744,28 @@ impl Assembler
                 Insn::Breakpoint => int3(cb),
 
                 Insn::CSelZ { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovnz);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovz, cmovnz);
                 },
                 Insn::CSelNZ { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovz);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovnz, cmovz);
                 },
                 Insn::CSelE { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovne);
+                    emit_csel(cb, *truthy, *falsy, *out, cmove, cmovne);
                 },
                 Insn::CSelNE { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmove);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovne, cmove);
                 },
                 Insn::CSelL { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovge);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovl, cmovge);
                 },
                 Insn::CSelLE { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovg);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovle, cmovg);
                 },
                 Insn::CSelG { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovle);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovg, cmovle);
                 },
                 Insn::CSelGE { truthy, falsy, out } => {
-                    emit_csel(cb, *truthy, *falsy, *out, cmovl);
+                    emit_csel(cb, *truthy, *falsy, *out, cmovge, cmovl);
                 }
                 Insn::LiveReg { .. } => (), // just a reg alloc signal, no code
                 Insn::PadInvalPatch => {
@@ -1177,4 +1197,26 @@ mod tests {
             0x23: call rax
         "});
     }
+
+    #[test]
+    fn test_cmov_mem() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let top = Opnd::mem(64, SP, 0);
+        let ary_opnd = SP;
+        let array_len_opnd = Opnd::mem(64, SP, 16);
+
+        asm.cmp(array_len_opnd, 1.into());
+        let elem_opnd = asm.csel_g(Opnd::mem(64, ary_opnd, 0), Qnil.into());
+        asm.mov(top, elem_opnd);
+
+        asm.compile_with_num_regs(&mut cb, 1);
+
+        assert_disasm!(cb, "48837b1001b804000000480f4f03488903", {"
+            0x0: cmp qword ptr [rbx + 0x10], 1
+            0x5: mov eax, 4
+            0xa: cmovg rax, qword ptr [rbx]
+            0xe: mov qword ptr [rbx], rax
+        "});
+    }
 }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index f8c52e58a8..ebcd5296fe 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1517,11 +1517,6 @@ fn gen_expandarray(
     let array_reg = asm.load(array_opnd);
     let array_len_opnd = get_array_len(asm, array_reg);
 
-    // Only handle the case where the number of values in the array is greater
-    // than or equal to the number of values requested.
-    asm.cmp(array_len_opnd, num.into());
-    asm.jl(Target::side_exit(Counter::expandarray_rhs_too_small));
-
     // Load the address of the embedded array into REG1.
     // (struct RArray *)(obj)->as.ary
     let array_reg = asm.load(array_opnd);
@@ -1542,7 +1537,11 @@ fn gen_expandarray(
     for i in (0..num).rev() {
         let top = asm.stack_push(Type::Unknown);
         let offset = i32::try_from(i * SIZEOF_VALUE).unwrap();
-        asm.mov(top, Opnd::mem(64, ary_opnd, offset));
+
+        // If the element index is less than the length of the array, load it
+        asm.cmp(array_len_opnd, i.into());
+        let elem_opnd = asm.csel_g(Opnd::mem(64, ary_opnd, offset), Qnil.into());
+        asm.mov(top, elem_opnd);
     }
 
     Some(KeepCompiling)