From 11c845efecc76735110e40d9b05ff5045f1ce925 Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Fri, 6 Feb 2026 21:43:23 -0500 Subject: ZJIT: Inline Primitives for Array#each --- array.c | 6 +++--- zjit/src/cruby_methods.rs | 42 ++++++++++++++++++++++++++++++++++++++++++ zjit/src/hir/opt_tests.rs | 25 +++++++++++++++++-------- zjit/src/hir/tests.rs | 12 ++++++------ 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 { + 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 { + 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 { + 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 , 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) "); -- cgit v1.2.3