From 90e6b01123ee129f37c01d69e99b32ffa0b44749 Mon Sep 17 00:00:00 2001
From: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Date: Wed, 2 Aug 2023 17:52:21 -0400
Subject: [PATCH] Propagate Type::Nil when known

---
 yjit/src/codegen.rs | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index fe933c5f5c36d9..dd0e690f78c530 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1470,7 +1470,7 @@ fn gen_expandarray(
     ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     // Both arguments are rb_num_t which is unsigned
-    let num = jit.get_arg(0).as_usize();
+    let num = jit.get_arg(0).as_u32();
     let flag = jit.get_arg(1).as_usize();
 
     // If this instruction has the splat flag, then bail out.
@@ -1535,7 +1535,34 @@ fn gen_expandarray(
     let array_reg = asm.load(array_opnd);
     let array_len_opnd = get_array_len(asm, array_reg);
 
+    /*
+    // FIXME: JCC_JB not implemented
     // Guard on the comptime/expected array length
+    if comptime_len >= num {
+        asm.comment(&format!("guard array length >= {}", num));
+        asm.cmp(array_len_opnd, num.into());
+        jit_chain_guard(
+            JCC_JB,
+            jit,
+            asm,
+            ocb,
+            OPT_AREF_MAX_CHAIN_DEPTH,
+            Counter::expandarray_chain_max_depth,
+        );
+    } else {
+        asm.comment(&format!("guard array length == {}", comptime_len));
+        asm.cmp(array_len_opnd, comptime_len.into());
+        jit_chain_guard(
+            JCC_JNE,
+            jit,
+            asm,
+            ocb,
+            OPT_AREF_MAX_CHAIN_DEPTH,
+            Counter::expandarray_chain_max_depth,
+        );
+    }
+    */
+
     asm.comment(&format!("guard array length == {}", comptime_len));
     asm.cmp(array_len_opnd, comptime_len.into());
     jit_chain_guard(
@@ -1575,12 +1602,12 @@ fn gen_expandarray(
 
         // Loop backward through the array and push each element onto the stack.
         for i in (0..num).rev() {
-            let top = asm.stack_push(Type::Unknown);
-            let offset = i32::try_from(i * SIZEOF_VALUE).unwrap();
+            let top = asm.stack_push(if i < comptime_len { Type::Unknown } else { Type::Nil });
+            let offset = i32::try_from(i * (SIZEOF_VALUE as u32)).unwrap();
 
             // Missing elements are Qnil
             asm.comment(&format!("load array[{}]", i));
-            let elem_opnd = if (i as u32) < comptime_len { Opnd::mem(64, ary_opnd, offset) } else { Qnil.into() };
+            let elem_opnd = if i < comptime_len { Opnd::mem(64, ary_opnd, offset) } else { Qnil.into() };
             asm.mov(top, elem_opnd);
         }
     }