From 3a9bf4a2ae1450bd8070296ed501f056830e25aa Mon Sep 17 00:00:00 2001 From: Max Bernstein Date: Mon, 23 Jun 2025 17:41:49 -0500 Subject: [PATCH] ZJIT: Optimize frozen array aref (#13666) If we have a frozen array `[..., a, ...]` and a compile-time fixnum index `i`, we can do the array load at compile-time. --- jit.c | 7 +++ yjit.c | 7 --- yjit/src/cruby_bindings.inc.rs | 2 +- zjit/src/cruby_bindings.inc.rs | 1 + zjit/src/hir.rs | 82 ++++++++++++++++++++++++++++++++++ 5 files changed, 91 insertions(+), 8 deletions(-) diff --git a/jit.c b/jit.c index d2147a9d7f..75ccd9b643 100644 --- a/jit.c +++ b/jit.c @@ -421,3 +421,10 @@ rb_assert_cme_handle(VALUE handle) RUBY_ASSERT_ALWAYS(!rb_objspace_garbage_object_p(handle)); RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_ment)); } + +// YJIT and ZJIT need this function to never allocate and never raise +VALUE +rb_yarv_ary_entry_internal(VALUE ary, long offset) +{ + return rb_ary_entry_internal(ary, offset); +} diff --git a/yjit.c b/yjit.c index ae042a62aa..ab527ef02f 100644 --- a/yjit.c +++ b/yjit.c @@ -533,13 +533,6 @@ rb_str_neq_internal(VALUE str1, VALUE str2) return rb_str_eql_internal(str1, str2) == Qtrue ? Qfalse : Qtrue; } -// YJIT needs this function to never allocate and never raise -VALUE -rb_yarv_ary_entry_internal(VALUE ary, long offset) -{ - return rb_ary_entry_internal(ary, offset); -} - extern VALUE rb_ary_unshift_m(int argc, VALUE *argv, VALUE ary); VALUE diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs index 8aa874f4dd..21ff8c7f06 100644 --- a/yjit/src/cruby_bindings.inc.rs +++ b/yjit/src/cruby_bindings.inc.rs @@ -1206,7 +1206,6 @@ extern "C" { pub fn rb_vm_base_ptr(cfp: *mut rb_control_frame_struct) -> *mut VALUE; pub fn rb_yarv_str_eql_internal(str1: VALUE, str2: VALUE) -> VALUE; pub fn rb_str_neq_internal(str1: VALUE, str2: VALUE) -> VALUE; - pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE; pub fn rb_ary_unshift_m(argc: ::std::os::raw::c_int, argv: *mut VALUE, ary: VALUE) -> VALUE; pub fn rb_yjit_rb_ary_subseq_length(ary: VALUE, beg: ::std::os::raw::c_long) -> VALUE; pub fn rb_yjit_fix_div_fix(recv: VALUE, obj: VALUE) -> VALUE; @@ -1328,4 +1327,5 @@ extern "C" { pub fn rb_assert_iseq_handle(handle: VALUE); pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int; pub fn rb_assert_cme_handle(handle: VALUE); + pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE; } diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs index d5e54955c8..518dc238ac 100644 --- a/zjit/src/cruby_bindings.inc.rs +++ b/zjit/src/cruby_bindings.inc.rs @@ -1001,4 +1001,5 @@ unsafe extern "C" { pub fn rb_assert_iseq_handle(handle: VALUE); pub fn rb_IMEMO_TYPE_P(imemo: VALUE, imemo_type: imemo_type) -> ::std::os::raw::c_int; pub fn rb_assert_cme_handle(handle: VALUE); + pub fn rb_yarv_ary_entry_internal(ary: VALUE, offset: ::std::os::raw::c_long) -> VALUE; } diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index f41a39e188..6ab58d7fb5 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -206,6 +206,7 @@ impl<'a> std::fmt::Display for InvariantPrinter<'a> { BOP_FREEZE => write!(f, "BOP_FREEZE")?, BOP_UMINUS => write!(f, "BOP_UMINUS")?, BOP_MAX => write!(f, "BOP_MAX")?, + BOP_AREF => write!(f, "BOP_AREF")?, _ => write!(f, "{bop}")?, } write!(f, ")") @@ -1317,6 +1318,25 @@ impl Function { } } + fn try_rewrite_aref(&mut self, block: BlockId, orig_insn_id: InsnId, self_val: InsnId, idx_val: InsnId) { + let self_type = self.type_of(self_val); + let idx_type = self.type_of(idx_val); + if self_type.is_subtype(types::ArrayExact) { + if let Some(array_obj) = self_type.ruby_object() { + if array_obj.is_frozen() { + if let Some(idx) = idx_type.fixnum_value() { + self.push_insn(block, Insn::PatchPoint(Invariant::BOPRedefined { klass: ARRAY_REDEFINED_OP_FLAG, bop: BOP_AREF })); + let val = unsafe { rb_yarv_ary_entry_internal(array_obj, idx) }; + let const_insn = self.push_insn(block, Insn::Const { val: Const::Value(val) }); + self.make_equal_to(orig_insn_id, const_insn); + return; + } + } + } + } + self.push_insn_id(block, orig_insn_id); + } + /// Rewrite SendWithoutBlock opcodes into SendWithoutBlockDirect opcodes if we know the target /// ISEQ statically. This removes run-time method lookups and opens the door for inlining. fn optimize_direct_sends(&mut self) { @@ -1351,6 +1371,8 @@ impl Function { self.try_rewrite_freeze(block, insn_id, self_val), Insn::SendWithoutBlock { self_val, call_info: CallInfo { method_name }, args, .. } if method_name == "-@" && args.len() == 0 => self.try_rewrite_uminus(block, insn_id, self_val), + Insn::SendWithoutBlock { self_val, call_info: CallInfo { method_name }, args, .. } if method_name == "[]" && args.len() == 1 => + self.try_rewrite_aref(block, insn_id, self_val, args[0]), Insn::SendWithoutBlock { mut self_val, call_info, cd, args, state } => { let frame_state = self.frame_state(state); let (klass, guard_equal_to) = if let Some(klass) = self.type_of(self_val).runtime_exact_ruby_class() { @@ -6068,4 +6090,64 @@ mod opt_tests { SideExit "#]]); } + + #[test] + fn test_eliminate_load_from_frozen_array_in_bounds() { + eval(r##" + def test = [4,5,6].freeze[1] + "##); + assert_optimized_method_hir("test", expect![[r#" + fn test: + bb0(v0:BasicObject): + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_FREEZE) + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_AREF) + v11:Fixnum[5] = Const Value(5) + Return v11 + "#]]); + } + + #[test] + fn test_eliminate_load_from_frozen_array_negative() { + eval(r##" + def test = [4,5,6].freeze[-3] + "##); + assert_optimized_method_hir("test", expect![[r#" + fn test: + bb0(v0:BasicObject): + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_FREEZE) + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_AREF) + v11:Fixnum[4] = Const Value(4) + Return v11 + "#]]); + } + + #[test] + fn test_eliminate_load_from_frozen_array_negative_out_of_bounds() { + eval(r##" + def test = [4,5,6].freeze[-10] + "##); + assert_optimized_method_hir("test", expect![[r#" + fn test: + bb0(v0:BasicObject): + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_FREEZE) + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_AREF) + v11:NilClassExact = Const Value(nil) + Return v11 + "#]]); + } + + #[test] + fn test_eliminate_load_from_frozen_array_out_of_bounds() { + eval(r##" + def test = [4,5,6].freeze[10] + "##); + assert_optimized_method_hir("test", expect![[r#" + fn test: + bb0(v0:BasicObject): + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_FREEZE) + PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_AREF) + v11:NilClassExact = Const Value(nil) + Return v11 + "#]]); + } }