ZJIT: Add Insn:: ArrayArefFixnum to accelerate Array#[] (#14717)

* ZJIT: Add Insn:: ArrayArefFixnum to accelerate Array#[]

* ZJIT: Use result from GuardType in ArrayArefFixnum

* ZJIT: Unbox index for aref_fixnum

* ZJIT: Change condition and add ArrayArefFixnum test

* ZJIT: Fix ArrayArefFixnum display for InsnPrinter

* ZJIT: Change insta test

5 files changed, 72 insertions, 0 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index c6dbc01dc2..faf717096a 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -1153,6 +1153,14 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_array_fixnum_aref
+    assert_compiles '3', %q{
+      def test(x) = [1,2,3][x]
+      test(2)
+      test(2)
+    }, call_threshold: 2, insns: [:opt_aref]
+  end
+
   def test_new_range_inclusive
     assert_compiles '1..5', %q{
       def test(a, b) = a..b
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index 64b235b838..b40986c0f6 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -125,6 +125,7 @@ fn main() {
         .allowlist_function("rb_ary_unshift_m")
         .allowlist_function("rb_ec_ary_new_from_values")
         .allowlist_function("rb_ary_tmp_new_from_values")
+        .allowlist_function("rb_ary_entry")
         .allowlist_function("rb_class_attached_object")
         .allowlist_function("rb_singleton_class")
         .allowlist_function("rb_define_class")
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 4b9331e05b..35791bc0d7 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -355,6 +355,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::NewRange { low, high, flag, state } => gen_new_range(jit, asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::NewRangeFixnum { low, high, flag, state } => gen_new_range_fixnum(asm, opnd!(low), opnd!(high), *flag, &function.frame_state(*state)),
         Insn::ArrayDup { val, state } => gen_array_dup(asm, opnd!(val), &function.frame_state(*state)),
+        Insn::ArrayArefFixnum { array, index, .. } => gen_aref_fixnum(asm, opnd!(array), opnd!(index)),
         Insn::ObjectAlloc { val, state } => gen_object_alloc(jit, asm, opnd!(val), &function.frame_state(*state)),
         &Insn::ObjectAllocClass { class, state } => gen_object_alloc_class(asm, class, &function.frame_state(state)),
         Insn::StringCopy { val, chilled, state } => gen_string_copy(asm, opnd!(val), *chilled, &function.frame_state(*state)),
@@ -1241,6 +1242,16 @@ fn gen_new_array(
     new_array
 }
 
+/// Compile array access (array[index])
+fn gen_aref_fixnum(
+    asm: &mut Assembler,
+    array: Opnd,
+    index: Opnd,
+) -> lir::Opnd {
+    let unboxed_idx = asm.rshift(index, Opnd::UImm(1));
+    asm_ccall!(asm, rb_ary_entry, array, unboxed_idx)
+}
+
 /// Compile a new hash instruction
 fn gen_new_hash(
     jit: &mut JITState,
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index ab442841ff..ea1bf68acc 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -786,6 +786,7 @@ unsafe extern "C" {
     pub fn rb_ary_resurrect(ary: VALUE) -> VALUE;
     pub fn rb_ary_cat(ary: VALUE, train: *const VALUE, len: ::std::os::raw::c_long) -> VALUE;
     pub fn rb_ary_push(ary: VALUE, elem: VALUE) -> VALUE;
+    pub fn rb_ary_entry(ary: VALUE, off: ::std::os::raw::c_long) -> VALUE;
     pub fn rb_ary_clear(ary: VALUE) -> VALUE;
     pub fn rb_ary_concat(lhs: VALUE, rhs: VALUE) -> VALUE;
     pub fn rb_hash_new() -> VALUE;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 2fe8eb7970..1ab6032832 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -577,6 +577,7 @@ pub enum Insn {
     ArrayExtend { left: InsnId, right: InsnId, state: InsnId },
     /// Push `val` onto `array`, where `array` is already `Array`.
     ArrayPush { array: InsnId, val: InsnId, state: InsnId },
+    ArrayArefFixnum { array: InsnId, index: InsnId },
 
     HashDup { val: InsnId, state: InsnId },
 
@@ -888,6 +889,10 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::ArrayArefFixnum { array, index, .. } => {
+                write!(f, "ArrayArefFixnum {array}, {index}")?;
+                Ok(())
+            }
             Insn::NewHash { elements, .. } => {
                 write!(f, "NewHash")?;
                 let mut prefix = " ";
@@ -1579,6 +1584,7 @@ impl Function {
             &NewHash { ref elements, state } => NewHash { elements: find_vec!(elements), state: find!(state) },
             &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state: find!(state) },
             &NewRangeFixnum { low, high, flag, state } => NewRangeFixnum { low: find!(low), high: find!(high), flag, state: find!(state) },
+            &ArrayArefFixnum { array, index } => ArrayArefFixnum { array: find!(array), index: find!(index) },
             &ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
             &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
@@ -1664,6 +1670,7 @@ impl Function {
             Insn::ToRegexp { .. } => types::RegexpExact,
             Insn::NewArray { .. } => types::ArrayExact,
             Insn::ArrayDup { .. } => types::ArrayExact,
+            Insn::ArrayArefFixnum { .. } => types::BasicObject,
             Insn::NewHash { .. } => types::HashExact,
             Insn::HashDup { .. } => types::HashExact,
             Insn::NewRange { .. } => types::RangeExact,
@@ -1969,6 +1976,16 @@ impl Function {
                     }
                 }
             }
+            if self.type_of(idx_val).is_subtype(types::Fixnum) {
+                self.push_insn(block, Insn::PatchPoint { invariant: Invariant::BOPRedefined { klass: ARRAY_REDEFINED_OP_FLAG, bop: BOP_AREF }, state });
+                let fixnum_idx = self.push_insn(block, Insn::GuardType { val: idx_val, guard_type: types::Fixnum, state });
+                let result = self.push_insn(block, Insn::ArrayArefFixnum {
+                    array: self_val,
+                    index: fixnum_idx,
+                });
+                self.make_equal_to(orig_insn_id, result);
+                return;
+            }
         }
         self.push_insn_id(block, orig_insn_id);
     }
@@ -2687,6 +2704,10 @@ impl Function {
                 worklist.push_back(val);
                 worklist.push_back(state);
             }
+            &Insn::ArrayArefFixnum { array, index } => {
+                worklist.push_back(array);
+                worklist.push_back(index);
+            }
             &Insn::Send { recv, ref args, state, .. }
             | &Insn::SendForward { recv, ref args, state, .. }
             | &Insn::SendWithoutBlock { recv, ref args, state, .. }
@@ -12564,4 +12585,34 @@ mod opt_tests {
           Return v23
         ");
     }
+
+    #[test]
+    fn test_array_aref_fixnum() {
+        eval("
+            def test
+              arr = [1, 2, 3]
+              arr[0]
+            end
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:NilClass = Const Value(nil)
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject):
+          EntryPoint JIT(0)
+          v6:NilClass = Const Value(nil)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:NilClass):
+          v13:ArrayExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v15:ArrayExact = ArrayDup v13
+          v18:Fixnum[0] = Const Value(0)
+          PatchPoint BOPRedefined(ARRAY_REDEFINED_OP_FLAG, BOP_AREF)
+          v30:BasicObject = ArrayArefFixnum v15, v18
+          CheckInterrupts
+          Return v30
+        ");
+    }
 }