ZJIT: Inline Primitives for Array#each

4 files changed, 68 insertions, 17 deletions

diff --git a/array.c b/array.c
index 876642dff7..c652c3c044 100644
--- a/array.c
+++ b/array.c
@@ -2695,21 +2695,21 @@ ary_enum_length(VALUE ary, VALUE args, VALUE eobj)
 }
 
 // Return true if the index is at or past the end of the array.
-static VALUE
+VALUE
 rb_jit_ary_at_end(rb_execution_context_t *ec, VALUE self, VALUE index)
 {
     return FIX2LONG(index) >= RARRAY_LEN(self) ? Qtrue : Qfalse;
 }
 
 // Return the element at the given fixnum index.
-static VALUE
+VALUE
 rb_jit_ary_at(rb_execution_context_t *ec, VALUE self, VALUE index)
 {
     return RARRAY_AREF(self, FIX2LONG(index));
 }
 
 // Increment a fixnum by 1.
-static VALUE
+VALUE
 rb_jit_fixnum_inc(rb_execution_context_t *ec, VALUE self, VALUE num)
 {
     return LONG2FIX(FIX2LONG(num) + 1);
diff --git a/zjit/src/cruby_methods.rs b/zjit/src/cruby_methods.rs
index aeca0e36bd..56ff3ca1aa 100644
--- a/zjit/src/cruby_methods.rs
+++ b/zjit/src/cruby_methods.rs
@@ -14,6 +14,13 @@ use std::ffi::c_void;
 use crate::hir_type::{types, Type};
 use crate::hir;
 
+// Array iteration builtin functions (defined in array.c)
+unsafe extern "C" {
+    fn rb_jit_ary_at_end(ec: EcPtr, self_: VALUE, index: VALUE) -> VALUE;
+    fn rb_jit_ary_at(ec: EcPtr, self_: VALUE, index: VALUE) -> VALUE;
+    fn rb_jit_fixnum_inc(ec: EcPtr, self_: VALUE, num: VALUE) -> VALUE;
+}
+
 pub struct Annotations {
     cfuncs: HashMap<*mut c_void, FnProperties>,
     builtin_funcs: HashMap<*mut c_void, FnProperties>,
@@ -270,6 +277,11 @@ pub fn init() -> Annotations {
     annotate_builtin!(rb_cSymbol, "name", types::StringExact);
     annotate_builtin!(rb_cSymbol, "to_s", types::StringExact);
 
+    // Array iteration builtins (used in with_jit Array#each, map, select, find)
+    builtin_funcs.insert(rb_jit_fixnum_inc as *mut c_void, FnProperties { inline: inline_fixnum_inc, return_type: types::Fixnum, ..Default::default() });
+    builtin_funcs.insert(rb_jit_ary_at as *mut c_void, FnProperties { inline: inline_ary_at, ..Default::default() });
+    builtin_funcs.insert(rb_jit_ary_at_end as *mut c_void, FnProperties { inline: inline_ary_at_end, return_type: types::BoolExact, ..Default::default() });
+
     Annotations {
         cfuncs: std::mem::take(cfuncs),
         builtin_funcs: std::mem::take(builtin_funcs),
@@ -892,3 +904,33 @@ fn inline_kernel_class(fun: &mut hir::Function, block: hir::BlockId, _recv: hir:
     let real_class = unsafe { rb_class_real(recv_class) };
     Some(fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(real_class) }))
 }
+
+/// Inline `fixnum_inc(ec, self, num)` implies FixnumAdd(num, 1).
+/// num is always a Fixnum (starts at 0 and is incremented by fixnum_inc).
+fn inline_fixnum_inc(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
+    let &[_self, num] = args else { return None; };
+    let one = fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(VALUE::fixnum_from_usize(1)) });
+    let result = fun.push_insn(block, hir::Insn::FixnumAdd { left: num, right: one, state });
+    Some(result)
+}
+
+/// Inline `ary_at(ec, self, index)` implies ArrayAref.
+/// Called from Array#each etc. where self is Array and index is in bounds.
+fn inline_ary_at(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
+    let &[recv, index] = args else { return None; };
+    let recv = fun.push_insn(block, hir::Insn::RefineType { val: recv, new_type: types::Array });
+    let index = fun.push_insn(block, hir::Insn::UnboxFixnum { val: index });
+    let result = fun.push_insn(block, hir::Insn::ArrayAref { array: recv, index });
+    Some(result)
+}
+
+/// Inline `ary_at_end(ec, self, index)` implies index >= ArrayLength(self).
+/// Called from Array#each etc. where self is Array and index is Fixnum.
+fn inline_ary_at_end(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
+    let &[recv, index] = args else { return None; };
+    let recv = fun.push_insn(block, hir::Insn::RefineType { val: recv, new_type: types::Array });
+    let length_cint = fun.push_insn(block, hir::Insn::ArrayLength { array: recv });
+    let length = fun.push_insn(block, hir::Insn::BoxFixnum { val: length_cint, state });
+    let result = fun.push_insn(block, hir::Insn::FixnumGe { left: index, right: length });
+    Some(result)
+}
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 8a172e6b12..d32de872d5 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -12030,18 +12030,27 @@ mod hir_opt_tests {
           v28:BasicObject = InvokeBuiltin <inline_expr>, v23
           CheckInterrupts
           Return v28
-        bb7(v48:BasicObject, v49:BasicObject):
-          v52:BasicObject = InvokeBuiltin rb_jit_ary_at_end, v48, v49
-          v54:CBool = Test v52
+        bb7(v48:BasicObject, v49:Fixnum):
+          v83:Array = RefineType v48, Array
+          v84:CInt64 = ArrayLength v83
+          v85:Fixnum = BoxFixnum v84
+          v86:BoolExact = FixnumGe v49, v85
+          IncrCounter inline_cfunc_optimized_send_count
+          v54:CBool = Test v86
           IfFalse v54, bb6(v48, v49)
           CheckInterrupts
           Return v48
-        bb6(v67:BasicObject, v68:BasicObject):
-          v72:BasicObject = InvokeBuiltin rb_jit_ary_at, v67, v68
-          v74:BasicObject = InvokeBlock, v72 # SendFallbackReason: Uncategorized(invokeblock)
-          v78:BasicObject = InvokeBuiltin rb_jit_fixnum_inc, v67, v68
+        bb6(v67:BasicObject, v68:Fixnum):
+          v88:Array = RefineType v67, Array
+          v89:CInt64 = UnboxFixnum v68
+          v90:BasicObject = ArrayAref v88, v89
+          IncrCounter inline_cfunc_optimized_send_count
+          v74:BasicObject = InvokeBlock, v90 # SendFallbackReason: Uncategorized(invokeblock)
+          v92:Fixnum[1] = Const Value(1)
+          v93:Fixnum = FixnumAdd v68, v92
+          IncrCounter inline_cfunc_optimized_send_count
           PatchPoint NoEPEscape(each)
-          Jump bb7(v67, v78)
+          Jump bb7(v67, v93)
         ");
     }
 }
diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs
index 8830b199c4..0401ffcd7d 100644
--- a/zjit/src/hir/tests.rs
+++ b/zjit/src/hir/tests.rs
@@ -4100,19 +4100,19 @@ pub mod hir_build_tests {
         bb5(v30:BasicObject, v31:NilClass):
           v35:Fixnum[0] = Const Value(0)
           Jump bb7(v30, v35)
-        bb7(v48:BasicObject, v49:BasicObject):
-          v52:BasicObject = InvokeBuiltin rb_jit_ary_at_end, v48, v49
+        bb7(v48:BasicObject, v49:Fixnum):
+          v52:BoolExact = InvokeBuiltin rb_jit_ary_at_end, v48, v49
           v54:CBool = Test v52
-          v55:Falsy = RefineType v52, Falsy
+          v55:FalseClass = RefineType v52, Falsy
           IfFalse v54, bb6(v48, v49)
-          v57:Truthy = RefineType v52, Truthy
+          v57:TrueClass = RefineType v52, Truthy
           v59:NilClass = Const Value(nil)
           CheckInterrupts
           Return v48
-        bb6(v67:BasicObject, v68:BasicObject):
+        bb6(v67:BasicObject, v68:Fixnum):
           v72:BasicObject = InvokeBuiltin rb_jit_ary_at, v67, v68
           v74:BasicObject = InvokeBlock, v72 # SendFallbackReason: Uncategorized(invokeblock)
-          v78:BasicObject = InvokeBuiltin rb_jit_fixnum_inc, v67, v68
+          v78:Fixnum = InvokeBuiltin rb_jit_fixnum_inc, v67, v68
           PatchPoint NoEPEscape(each)
           Jump bb7(v67, v78)
         ");