ZJIT: Strength-reduce Integer#[] with a constant index

Author: k0kubun

zjit/src/codegen_tests.rs

@@ -2842,6 +2842,38 @@ fn test_fixnum_mod_negative() {
     assert_snapshot!(assert_compiles("[test(-7, 3), test(7, -3), test(-7, -3)]"), @"[2, -2, -1]");
 }
 
+#[test]
+fn test_fixnum_aref_constant_index() {
+    eval("
+        def test(a) = a[12]
+        test(4096) # profile opt_aref
+    ");
+    assert_contains_opcode("test", YARVINSN_opt_aref);
+    assert_snapshot!(assert_compiles("[test(4096), test(4095), test(0), test(-1), test(-4096)]"), @"[1, 0, 0, 1, 1]");
+}
+
+#[test]
+fn test_fixnum_aref_constant_index_beyond_fixnum_width() {
+    // An index beyond the fixnum width is not strength-reduced; FixnumAref handles it
+    eval("
+        def test(a) = a[100]
+        test(1) # profile opt_aref
+    ");
+    assert_contains_opcode("test", YARVINSN_opt_aref);
+    assert_snapshot!(assert_compiles("[test(1), test(-1), test(4611686018427387903), test(-4611686018427387904)]"), @"[0, 1, 0, 1]");
+}
+
+#[test]
+fn test_fixnum_aref_constant_index_bignum_receiver() {
+    // A Bignum receiver fails the Fixnum guard and side-exits to the correct result
+    eval("
+        def test(a) = a[1]
+        test(5) # profile opt_aref
+    ");
+    assert_contains_opcode("test", YARVINSN_opt_aref);
+    assert_snapshot!(assert_compiles_allowing_exits("[test(5), test(2**100 + 2)]"), @"[0, 1]");
+}
+
 #[test]
 fn test_fixnum_mod_by_zero() {
     eval("

zjit/src/cruby_methods.rs

@@ -800,6 +800,24 @@ fn inline_integer_rshift(fun: &mut hir::Function, block: hir::BlockId, recv: hir
 
 fn inline_integer_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
     let &[index] = args else { return None; };
+
+    // Optimize a compile-time constant index that fits in the Fixnum payload as a bit test.
+    // A Fixnum payload is VALUE_BITS - 1 bits wide (one bit is the tag), so its highest
+    // (sign) bit is at index VALUE_BITS - 2.
+    const FIXNUM_SIGN_BIT_INDEX: i64 = VALUE_BITS as i64 - 2;
+    if fun.likely_a(recv, types::Fixnum, state) {
+        if let Some(index_value) = fun.type_of(index).fixnum_value() {
+            if (0..=FIXNUM_SIGN_BIT_INDEX).contains(&index_value) {
+                // Optimize `recv[i]` into `(recv >> i) & 1`.
+                let recv = fun.coerce_to(block, recv, types::Fixnum, state);
+                let shift = fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(VALUE::fixnum_from_usize(index_value as usize)) });
+                let shifted = fun.push_insn(block, hir::Insn::FixnumRShift { left: recv, right: shift });
+                let mask = fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(VALUE::fixnum_from_usize(1)) });
+                return Some(fun.push_insn(block, hir::Insn::FixnumAnd { left: shifted, right: mask }));
+            }
+        }
+    }
+    // Use a C call (FixnumAref) for any other (e.g. non-constant) index.
     if fun.likely_a(recv, types::Fixnum, state) && fun.likely_a(index, types::Fixnum, state) {
         let recv = fun.coerce_to(block, recv, types::Fixnum, state);
         let index = fun.coerce_to(block, index, types::Fixnum, state);

zjit/src/hir/opt_tests.rs

@@ -2064,6 +2064,67 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn integer_aref_with_constant_index_strength_reduced() {
+        eval("
+            def test(a) = a[12]
+            test(4096)
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:CPtr = LoadSP
+          v3:BasicObject = LoadField v2, :a@0x1000
+          Jump bb3(v1, v3)
+        bb2():
+          EntryPoint JIT(0)
+          v6:BasicObject = LoadArg :self@0
+          v7:BasicObject = LoadArg :a@1
+          Jump bb3(v6, v7)
+        bb3(v9:BasicObject, v10:BasicObject):
+          v15:Fixnum[12] = Const Value(12)
+          PatchPoint MethodRedefined(Integer@0x1008, []@0x1010, cme:0x1018)
+          v26:Fixnum = GuardType v10, Fixnum recompile
+          v27:Fixnum[12] = Const Value(12)
+          v28:Fixnum = FixnumRShift v26, v27
+          v29:Fixnum[1] = Const Value(1)
+          v30:Fixnum = FixnumAnd v28, v29
+          CheckInterrupts
+          Return v30
+        ");
+    }
+
+    #[test]
+    fn integer_aref_with_constant_index_beyond_fixnum_width() {
+        eval("
+            def test(a) = a[100]
+            test(-1)
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:CPtr = LoadSP
+          v3:BasicObject = LoadField v2, :a@0x1000
+          Jump bb3(v1, v3)
+        bb2():
+          EntryPoint JIT(0)
+          v6:BasicObject = LoadArg :self@0
+          v7:BasicObject = LoadArg :a@1
+          Jump bb3(v6, v7)
+        bb3(v9:BasicObject, v10:BasicObject):
+          v15:Fixnum[100] = Const Value(100)
+          PatchPoint MethodRedefined(Integer@0x1008, []@0x1010, cme:0x1018)
+          v26:Fixnum = GuardType v10, Fixnum recompile
+          v27:Fixnum = FixnumAref v26, v15
+          CheckInterrupts
+          Return v27
+        ");
+    }
+
     #[test]
     fn elide_fixnum_aref() {
         eval("