ZJIT: Inline ArrayAref

Merged: https://github.com/ruby/ruby/pull/15833

test/ruby/test_zjit.rb

@@ -2160,6 +2160,54 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 2, insns: [:opt_aref]
   end
 
+  def test_array_fixnum_aref_negative_index
+    assert_compiles '3', %q{
+      def test(x) = [1,2,3][x]
+      test(-1)
+      test(-1)
+    }, call_threshold: 2, insns: [:opt_aref]
+  end
+
+  def test_array_fixnum_aref_out_of_bounds_positive
+    assert_compiles 'nil', %q{
+      def test(x) = [1,2,3][x]
+      test(10)
+      test(10)
+    }, call_threshold: 2, insns: [:opt_aref]
+  end
+
+  def test_array_fixnum_aref_out_of_bounds_negative
+    assert_compiles 'nil', %q{
+      def test(x) = [1,2,3][x]
+      test(-10)
+      test(-10)
+    }, call_threshold: 2, insns: [:opt_aref]
+  end
+
+  def test_array_fixnum_aref_array_subclass
+    assert_compiles '3', %q{
+      class MyArray < Array; end
+      def test(arr, idx) = arr[idx]
+      arr = MyArray[1,2,3]
+      test(arr, 2)
+      arr = MyArray[1,2,3]
+      test(arr, 2)
+    }, call_threshold: 2, insns: [:opt_aref]
+  end
+
+  def test_array_aref_non_fixnum_index
+    assert_compiles 'TypeError', %q{
+      def test(arr, idx) = arr[idx]
+      test([1,2,3], 1)
+      test([1,2,3], 1)
+      begin
+        test([1,2,3], "1")
+      rescue => e
+        e.class
+      end
+    }, call_threshold: 2
+  end
+
   def test_array_fixnum_aset
     assert_compiles '[1, 2, 7]', %q{
       def test(arr, idx)

zjit/src/codegen.rs

@@ -1565,7 +1565,12 @@ fn gen_array_aref(
     array: Opnd,
     index: Opnd,
 ) -> lir::Opnd {
-    asm_ccall!(asm, rb_ary_entry, array, index)
+    let unboxed_idx = asm.load(index);
+    let array = asm.load(array);
+    let array_ptr = gen_array_ptr(asm, array);
+    let elem_offset = asm.lshift(unboxed_idx, Opnd::UImm(SIZEOF_VALUE.trailing_zeros() as u64));
+    let elem_ptr = asm.add(array_ptr, elem_offset);
+    asm.load(Opnd::mem(VALUE_BITS, elem_ptr, 0))
 }
 
 fn gen_array_aset(

zjit/src/cruby_methods.rs

@@ -327,6 +327,10 @@ fn inline_array_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::In
         if fun.likely_a(index, types::Fixnum, state) {
             let index = fun.coerce_to(block, index, types::Fixnum, state);
             let index = fun.push_insn(block, hir::Insn::UnboxFixnum { val: index });
+            let length = fun.push_insn(block, hir::Insn::ArrayLength { array: recv });
+            let index = fun.push_insn(block, hir::Insn::GuardLess { left: index, right: length, state });
+            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, state });
             let result = fun.push_insn(block, hir::Insn::ArrayAref { array: recv, index });
             return Some(result);
         }

zjit/src/hir.rs

@@ -2251,13 +2251,10 @@ impl Function {
             Insn::IsBitNotEqual { .. } => types::CBool,
             Insn::BoxBool { .. } => types::BoolExact,
             Insn::BoxFixnum { .. } => types::Fixnum,
-            Insn::UnboxFixnum { val } => {
-                if let Some(fixnum) = self.type_of(*val).fixnum_value() {
-                    Type::from_cint(types::CInt64, fixnum)
-                } else {
-                    types::CInt64
-                }
-            },
+            Insn::UnboxFixnum { val } => self
+                .type_of(*val)
+                .fixnum_value()
+                .map_or(types::CInt64, |fixnum| Type::from_cint(types::CInt64, fixnum)),
             Insn::FixnumAref { .. } => types::Fixnum,
             Insn::StringCopy { .. } => types::StringExact,
             Insn::StringIntern { .. } => types::Symbol,
@@ -6548,9 +6545,9 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let length = fun.push_insn(block, Insn::ArrayLength { array });
                     fun.push_insn(block, Insn::GuardBitEquals { val: length, expected: Const::CInt64(num as i64), reason: SideExitReason::ExpandArray, state: exit_id });
                     for i in (0..num).rev() {
-                        // TODO(max): Add a short-cut path for long indices into an array where the
-                        // index is known to be in-bounds
                         let index = fun.push_insn(block, Insn::Const { val: Const::Value(VALUE::fixnum_from_usize(i.try_into().unwrap())) });
+                        // We do not emit a length guard here because in-bounds is already
+                        // ensured by the expandarray length check above.
                         let index = fun.push_insn(block, Insn::UnboxFixnum { val: index });
                         let element = fun.push_insn(block, Insn::ArrayAref { array, index });
                         state.stack_push(element);

zjit/src/hir/opt_tests.rs

@@ -1607,10 +1607,14 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Array@0x1000, []@0x1008, cme:0x1010)
           v25:ArrayExact = GuardType v9, ArrayExact
           v26:CInt64[0] = UnboxFixnum v14
-          v27:BasicObject = ArrayAref v25, v26
+          v27:CInt64 = ArrayLength v25
+          v28:CInt64[0] = GuardLess v26, v27
+          v29:CInt64[0] = Const CInt64(0)
+          v30:CInt64[0] = GuardGreaterEq v28, v29
+          v31:BasicObject = ArrayAref v25, v30
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v27
+          Return v31
         ");
         assert_snapshot!(inspect("test [1,2,3]"), @"1");
     }
@@ -4736,10 +4740,14 @@ mod hir_opt_tests {
           PatchPoint NoSingletonClass(Array@0x1010)
           PatchPoint MethodRedefined(Array@0x1010, []@0x1018, cme:0x1020)
           v27:CInt64[0] = UnboxFixnum v13
-          v28:BasicObject = ArrayAref v23, v27
+          v28:CInt64 = ArrayLength v23
+          v29:CInt64[0] = GuardLess v27, v28
+          v30:CInt64[0] = Const CInt64(0)
+          v31:CInt64[0] = GuardGreaterEq v29, v30
+          v32:BasicObject = ArrayAref v23, v31
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v28
+          Return v32
         ");
        // TODO(max): Check the result of `S[0] = 5; test` using `inspect` to make sure that we
        // actually do the load at run-time.
@@ -4766,10 +4774,14 @@ mod hir_opt_tests {
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, []@0x1010, cme:0x1018)
           v24:CInt64[1] = UnboxFixnum v13
-          v27:Fixnum[5] = Const Value(5)
+          v25:CInt64 = ArrayLength v11
+          v26:CInt64[1] = GuardLess v24, v25
+          v27:CInt64[0] = Const CInt64(0)
+          v28:CInt64[1] = GuardGreaterEq v26, v27
+          v31:Fixnum[5] = Const Value(5)
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v27
+          Return v31
         ");
     }
 
@@ -4794,10 +4806,14 @@ mod hir_opt_tests {
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, []@0x1010, cme:0x1018)
           v24:CInt64[-3] = UnboxFixnum v13
-          v27:Fixnum[4] = Const Value(4)
+          v25:CInt64 = ArrayLength v11
+          v26:CInt64[-3] = GuardLess v24, v25
+          v27:CInt64[0] = Const CInt64(0)
+          v28:CInt64[-3] = GuardGreaterEq v26, v27
+          v31:Fixnum[4] = Const Value(4)
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v27
+          Return v31
         ");
     }
 
@@ -4822,10 +4838,14 @@ mod hir_opt_tests {
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, []@0x1010, cme:0x1018)
           v24:CInt64[-10] = UnboxFixnum v13
-          v27:NilClass = Const Value(nil)
+          v25:CInt64 = ArrayLength v11
+          v26:CInt64[-10] = GuardLess v24, v25
+          v27:CInt64[0] = Const CInt64(0)
+          v28:CInt64[-10] = GuardGreaterEq v26, v27
+          v31:NilClass = Const Value(nil)
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v27
+          Return v31
         ");
     }
 
@@ -4850,10 +4870,14 @@ mod hir_opt_tests {
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, []@0x1010, cme:0x1018)
           v24:CInt64[10] = UnboxFixnum v13
-          v27:NilClass = Const Value(nil)
+          v25:CInt64 = ArrayLength v11
+          v26:CInt64[10] = GuardLess v24, v25
+          v27:CInt64[0] = Const CInt64(0)
+          v28:CInt64[10] = GuardGreaterEq v26, v27
+          v31:NilClass = Const Value(nil)
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v27
+          Return v31
         ");
     }
 
@@ -6742,10 +6766,14 @@ mod hir_opt_tests {
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, []@0x1010, cme:0x1018)
           v30:CInt64[0] = UnboxFixnum v19
-          v31:BasicObject = ArrayAref v14, v30
+          v31:CInt64 = ArrayLength v14
+          v32:CInt64[0] = GuardLess v30, v31
+          v33:CInt64[0] = Const CInt64(0)
+          v34:CInt64[0] = GuardGreaterEq v32, v33
+          v35:BasicObject = ArrayAref v14, v34
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v31
+          Return v35
         ");
     }
 
@@ -6774,10 +6802,14 @@ mod hir_opt_tests {
           v27:ArrayExact = GuardType v11, ArrayExact
           v28:Fixnum = GuardType v12, Fixnum
           v29:CInt64 = UnboxFixnum v28
-          v30:BasicObject = ArrayAref v27, v29
+          v30:CInt64 = ArrayLength v27
+          v31:CInt64 = GuardLess v29, v30
+          v32:CInt64[0] = Const CInt64(0)
+          v33:CInt64 = GuardGreaterEq v31, v32
+          v34:BasicObject = ArrayAref v27, v33
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v30
+          Return v34
         ");
     }
 
@@ -6807,10 +6839,14 @@ mod hir_opt_tests {
           v27:ArraySubclass[class_exact:C] = GuardType v11, ArraySubclass[class_exact:C]
           v28:Fixnum = GuardType v12, Fixnum
           v29:CInt64 = UnboxFixnum v28
-          v30:BasicObject = ArrayAref v27, v29
+          v30:CInt64 = ArrayLength v27
+          v31:CInt64 = GuardLess v29, v30
+          v32:CInt64[0] = Const CInt64(0)
+          v33:CInt64 = GuardGreaterEq v31, v32
+          v34:BasicObject = ArrayAref v27, v33
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
-          Return v30
+          Return v34
         ");
     }
 

zjit/src/hir_type/mod.rs

@@ -396,13 +396,9 @@ impl Type {
     }
 
     pub fn cint64_value(&self) -> Option<i64> {
-        if self.is_subtype(types::CInt64) {
-            match self.spec {
-                Specialization::Int(val) => Some(val as i64),
-                _ => None,
-            }
-        } else {
-            None
+        match (self.is_subtype(types::CInt64), &self.spec) {
+            (true, Specialization::Int(val)) => Some(*val as i64),
+            _ => None,
         }
     }