ZJIT: Inline Array#first and Array#last

Author: k0kubun

array.c

@@ -2730,6 +2730,20 @@ rb_builtin_fixnum_inc(rb_execution_context_t *ec, VALUE self, VALUE num)
     return LONG2FIX(FIX2LONG(num) + 1);
 }
 
+// Return the first element of the array, or nil if empty.
+VALUE
+rb_builtin_ary_first(rb_execution_context_t *ec, VALUE self)
+{
+    return ary_first(self);
+}
+
+// Return the last element of the array, or nil if empty.
+VALUE
+rb_builtin_ary_last(rb_execution_context_t *ec, VALUE self)
+{
+    return ary_last(self);
+}
+
 // Push a value onto an array and return the value.
 static VALUE
 rb_jit_ary_push(rb_execution_context_t *ec, VALUE self, VALUE ary, VALUE val)

array.rb

@@ -129,10 +129,10 @@ def sample(n = (ary = false), random: Random)
   def first n = unspecified = true
     if Primitive.mandatory_only?
       Primitive.attr! :leaf
-      Primitive.cexpr! %q{ ary_first(self) }
+      Primitive.rb_builtin_ary_first
     else
       if unspecified
-        Primitive.cexpr! %q{ ary_first(self) }
+        Primitive.rb_builtin_ary_first
       else
         Primitive.cexpr! %q{  ary_take_first_or_last_n(self, NUM2LONG(n), ARY_TAKE_FIRST) }
       end
@@ -166,10 +166,10 @@ def first n = unspecified = true
   def last n = unspecified = true
     if Primitive.mandatory_only?
       Primitive.attr! :leaf
-      Primitive.cexpr! %q{ ary_last(self) }
+      Primitive.rb_builtin_ary_last
     else
       if unspecified
-        Primitive.cexpr! %q{ ary_last(self) }
+        Primitive.rb_builtin_ary_last
       else
         Primitive.cexpr! %q{ ary_take_first_or_last_n(self, NUM2LONG(n), ARY_TAKE_LAST) }
       end

zjit/src/cruby_methods.rs

@@ -19,6 +19,8 @@ unsafe extern "C" {
     fn rb_builtin_ary_at_end(ec: EcPtr, self_: VALUE, index: VALUE) -> VALUE;
     fn rb_builtin_ary_at(ec: EcPtr, self_: VALUE, index: VALUE) -> VALUE;
     fn rb_builtin_fixnum_inc(ec: EcPtr, self_: VALUE, num: VALUE) -> VALUE;
+    fn rb_builtin_ary_first(ec: EcPtr, self_: VALUE) -> VALUE;
+    fn rb_builtin_ary_last(ec: EcPtr, self_: VALUE) -> VALUE;
     fn rb_str_equal(str1: VALUE, str2: VALUE) -> VALUE;
 }
 
@@ -301,6 +303,10 @@ pub fn init() -> Annotations {
     builtin_funcs.insert(rb_builtin_ary_at as *mut c_void, FnProperties { inline: inline_ary_at, ..Default::default() });
     builtin_funcs.insert(rb_builtin_ary_at_end as *mut c_void, FnProperties { inline: inline_ary_at_end, return_type: types::BoolExact, ..Default::default() });
 
+    // Builtins used by Array#first and Array#last when called with no arguments
+    builtin_funcs.insert(rb_builtin_ary_first as *mut c_void, FnProperties { inline: inline_ary_first, no_gc: true, leaf: true, elidable: true, ..Default::default() });
+    builtin_funcs.insert(rb_builtin_ary_last as *mut c_void, FnProperties { inline: inline_ary_last, no_gc: true, leaf: true, elidable: true, ..Default::default() });
+
     Annotations {
         cfuncs: std::mem::take(cfuncs),
         builtin_funcs: std::mem::take(builtin_funcs),
@@ -1057,3 +1063,35 @@ fn inline_ary_at_end(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::I
     let result = fun.push_insn(block, hir::Insn::FixnumGe { left: index, right: length });
     Some(result)
 }
+
+/// Inline `ary_first(ec, self)` from Array#first with no arguments.
+/// Guards that the array is non-empty and loads the element at index 0;
+/// side-exits on empty arrays and lets the interpreter return nil.
+fn inline_ary_first(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
+    use crate::hir::SideExitReason;
+    let &[recv] = args else { return None; };
+    let recv = fun.push_insn(block, hir::Insn::RefineType { val: recv, new_type: types::Array });
+    let length = fun.push_insn(block, hir::Insn::ArrayLength { array: recv });
+    let zero = fun.push_insn(block, hir::Insn::Const { val: hir::Const::CInt64(0) });
+    let index = fun.push_insn(block, hir::Insn::GuardLess { left: zero, right: length, reason: Box::new(SideExitReason::GuardLess), state });
+    let result = fun.push_insn(block, hir::Insn::ArrayAref { array: recv, index });
+    Some(result)
+}
+
+/// Inline `ary_last(ec, self)` from Array#last with no arguments.
+/// Guards that the array is non-empty and loads the element at index length-1;
+/// side-exits on empty arrays and lets the interpreter return nil.
+fn inline_ary_last(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
+    use crate::hir::SideExitReason;
+    let &[recv] = args else { return None; };
+    let recv = fun.push_insn(block, hir::Insn::RefineType { val: recv, new_type: types::Array });
+    let length = fun.push_insn(block, hir::Insn::ArrayLength { array: recv });
+    let minus_one = fun.push_insn(block, hir::Insn::Const { val: hir::Const::CInt64(-1) });
+    // AdjustBounds computes length - 1 since the index is negative
+    // TODO: Consider doing something like IntSub instead
+    let index = fun.push_insn(block, hir::Insn::AdjustBounds { index: minus_one, length });
+    let zero = fun.push_insn(block, hir::Insn::Const { val: hir::Const::CInt64(0) });
+    let index = fun.push_insn(block, hir::Insn::GuardGreaterEq { left: index, right: zero, reason: Box::new(SideExitReason::GuardGreaterEq), state });
+    let result = fun.push_insn(block, hir::Insn::ArrayAref { array: recv, index });
+    Some(result)
+}

zjit/src/hir/opt_tests.rs

@@ -3688,9 +3688,12 @@ mod hir_opt_tests {
           v18:ArrayExact = ArrayDup v17
           PatchPoint NoSingletonClass(Array@0x1010)
           PatchPoint MethodRedefined(Array@0x1010, first@0x1018, cme:0x1020)
-          v31:BasicObject = InvokeBuiltin leaf <inline_expr>, v18
+          v32:CInt64 = ArrayLength v18
+          v33:CInt64[0] = Const CInt64(0)
+          v34:CInt64[0] = GuardLess v33, v32
+          v35:BasicObject = ArrayAref v18, v34
           CheckInterrupts
-          Return v31
+          Return v35
         ");
     }
 
@@ -11782,6 +11785,72 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_optimize_array_first() {
+        eval("
+            def test(arr) = arr.first
+            test([1])
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:CPtr = LoadSP
+          v3:BasicObject = LoadField v2, :arr@0x1000
+          Jump bb3(v1, v3)
+        bb2():
+          EntryPoint JIT(0)
+          v6:BasicObject = LoadArg :self@0
+          v7:BasicObject = LoadArg :arr@1
+          Jump bb3(v6, v7)
+        bb3(v9:BasicObject, v10:BasicObject):
+          PatchPoint NoSingletonClass(Array@0x1008)
+          PatchPoint MethodRedefined(Array@0x1008, first@0x1010, cme:0x1018)
+          v23:ArrayExact = GuardType v10, ArrayExact recompile
+          v25:CInt64 = ArrayLength v23
+          v26:CInt64[0] = Const CInt64(0)
+          v27:CInt64[0] = GuardLess v26, v25
+          v28:BasicObject = ArrayAref v23, v27
+          CheckInterrupts
+          Return v28
+        ");
+    }
+
+    #[test]
+    fn test_optimize_array_last() {
+        eval("
+            def test(arr) = arr.last
+            test([1])
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:CPtr = LoadSP
+          v3:BasicObject = LoadField v2, :arr@0x1000
+          Jump bb3(v1, v3)
+        bb2():
+          EntryPoint JIT(0)
+          v6:BasicObject = LoadArg :self@0
+          v7:BasicObject = LoadArg :arr@1
+          Jump bb3(v6, v7)
+        bb3(v9:BasicObject, v10:BasicObject):
+          PatchPoint NoSingletonClass(Array@0x1008)
+          PatchPoint MethodRedefined(Array@0x1008, last@0x1010, cme:0x1018)
+          v23:ArrayExact = GuardType v10, ArrayExact recompile
+          v25:CInt64 = ArrayLength v23
+          v26:CInt64[-1] = Const CInt64(-1)
+          v27:CInt64 = AdjustBounds v26, v25
+          v28:CInt64[0] = Const CInt64(0)
+          v29:CInt64 = GuardGreaterEq v27, v28
+          v30:BasicObject = ArrayAref v23, v29
+          CheckInterrupts
+          Return v30
+        ");
+    }
+
     #[test]
     fn test_optimize_regexpmatch2() {
         eval(r#"