ZJIT: Optimize Array#[]= with splice semantics (index, length)

Add inlined implementation for the three-argument form of Array#[]=
(arr[index, len] = val) which splices elements into the array. Guards
for exact Array type, Fixnum index and length, and that the array is
not frozen or shared. Falls back to the standard method for Array
subclasses.

Author: nmeans

array.c

@@ -2715,6 +2715,13 @@ rb_jit_fixnum_inc(rb_execution_context_t *ec, VALUE self, VALUE num)
     return LONG2FIX(FIX2LONG(num) + 1);
 }
 
+// Splice a value into an array at the given index and length.
+VALUE
+rb_jit_ary_aset_by_rb_ary_splice(VALUE ary, long beg, long len, VALUE val)
+{
+    return ary_aset_by_rb_ary_splice(ary, beg, len, val);
+}
+
 // Push a value onto an array and return the value.
 VALUE
 rb_jit_ary_push(rb_execution_context_t *ec, VALUE self, VALUE ary, VALUE val)

jit.c

@@ -791,6 +791,7 @@ rb_yarv_str_eql_internal(VALUE str1, VALUE str2)
 }
 
 void rb_jit_str_concat_codepoint(VALUE str, VALUE codepoint);
+VALUE rb_jit_ary_aset_by_rb_ary_splice(VALUE ary, long beg, long len, VALUE val);
 
 attr_index_t
 rb_jit_shape_capacity(shape_id_t shape_id)

zjit/bindgen/src/main.rs

@@ -283,6 +283,7 @@ fn main() {
         .allowlist_function("rb_jit_get_page_size")
         .allowlist_function("rb_jit_array_len")
         .allowlist_function("rb_jit_fix_div_fix")
+        .allowlist_function("rb_jit_ary_aset_by_rb_ary_splice")
         .allowlist_function("rb_jit_iseq_builtin_attrs")
         .allowlist_function("rb_jit_str_concat_codepoint")
         .allowlist_function("rb_zjit_iseq_inspect")

zjit/src/codegen_tests.rs

@@ -2781,6 +2781,72 @@ fn test_array_fixnum_aset_array_subclass() {
     assert_snapshot!(inspect("arr = MyArray.new; test(arr, 0); arr[0]"), @"7");
 }
 
+#[test]
+fn test_array_splice_aset() {
+    eval("
+        def test(arr, idx, len)
+            arr[idx, len] = [4, 5, 6]
+        end
+        test([1,2,3], 0, 2)
+    ");
+    assert_snapshot!(inspect("arr = [1,2,3]; test(arr, 0, 2); arr"), @"[4, 5, 6, 3]");
+}
+
+#[test]
+fn test_array_splice_aset_returns_value() {
+    eval("
+        def test(arr, idx, len, val)
+            arr[idx, len] = val
+        end
+        test([1,2,3], 0, 2, [4,5,6])
+    ");
+    assert_snapshot!(inspect("test([1,2,3], 0, 2, [4,5,6])"), @"[4, 5, 6]");
+}
+
+#[test]
+fn test_array_splice_aset_frozen() {
+    assert_snapshot!(inspect("
+        def test(arr, idx, len, val)
+            arr[idx, len] = val
+        end
+        arr = [1,2,3]
+        test(arr, 0, 2, [4,5])
+        test(arr, 0, 2, [4,5])
+        arr.freeze
+        begin
+            test(arr, 0, 2, [4,5])
+        rescue => e
+            e.class
+        end
+    "), @"FrozenError");
+}
+
+#[test]
+fn test_array_splice_aset_shared() {
+    assert_snapshot!(inspect("
+        def test(arr, idx, len, val)
+            arr[idx, len] = val
+        end
+        arr = (0..50).to_a
+        test(arr, 0, 1, [-1])
+        shared = arr[10, 20]
+        test(shared, 0, 1, [999])
+        [arr[10], shared[0], arr[0]]
+    "), @"[10, 999, -1]");
+}
+
+#[test]
+fn test_array_splice_aset_array_subclass() {
+    eval("
+        class MySpliceArray < Array; end
+        def test(arr, idx, len, val)
+            arr[idx, len] = val
+        end
+        test(MySpliceArray.new([1,2,3]), 0, 2, [4,5])
+    ");
+    assert_snapshot!(inspect("arr = MySpliceArray.new([1,2,3]); test(arr, 0, 2, [4,5]); arr.to_a"), @"[4, 5, 3]");
+}
+
 #[test]
 fn test_array_aset_non_fixnum_index() {
     assert_snapshot!(inspect(r#"

zjit/src/cruby.rs

@@ -1416,6 +1416,7 @@ pub(crate) mod ids {
         name: freeze
         name: minusat            content: b"-@"
         name: aref               content: b"[]"
+        name: aset               content: b"[]="
         name: len
         name: _as_heap
         name: thread_ptr

zjit/src/cruby_bindings.inc.rs

@@ -379,6 +379,29 @@ fn inline_array_aset(fun: &mut hir::Function, block: hir::BlockId, recv: hir::In
             fun.push_insn(block, hir::Insn::WriteBarrier { recv, val });
             return Some(val);
         }
+    } else if let &[index, len, val] = args {
+        if fun.likely_a(recv, types::ArrayExact, state)
+            && fun.likely_a(index, types::Fixnum, state)
+            && fun.likely_a(len, types::Fixnum, state)
+        {
+            let recv = fun.coerce_to(block, recv, types::ArrayExact, state);
+            let index = fun.coerce_to(block, index, types::Fixnum, state);
+            let len = fun.coerce_to(block, len, types::Fixnum, state);
+            fun.guard_not_frozen(block, recv, state);
+            fun.guard_not_shared(block, recv, state);
+
+            let index = fun.push_insn(block, hir::Insn::UnboxFixnum { val: index });
+            let len = fun.push_insn(block, hir::Insn::UnboxFixnum { val: len });
+            let _ = fun.push_insn(block, hir::Insn::CCall {
+                cfunc: rb_jit_ary_aset_by_rb_ary_splice as _,
+                recv,
+                args: vec![index, len, val],
+                return_type: types::BasicObject,
+                elidable: false,
+                name: ID!(aset),
+            });
+            return Some(val);
+        }
     }
     None
 }

zjit/src/cruby_methods.rs

@@ -8383,6 +8383,89 @@ mod hir_opt_tests {
         ");
     }
 
+    #[test]
+    fn test_optimize_array_splice_aset_profiled() {
+        eval("
+            def test(arr, index, len, val)
+              arr[index, len] = val
+            end
+            test([1,2,3], 0, 2, [4,5,6])
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:3:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal :arr, l0, SP@7
+          v3:BasicObject = GetLocal :index, l0, SP@6
+          v4:BasicObject = GetLocal :len, l0, SP@5
+          v5:BasicObject = GetLocal :val, l0, SP@4
+          Jump bb3(v1, v2, v3, v4, v5)
+        bb2():
+          EntryPoint JIT(0)
+          v8:BasicObject = LoadArg :self@0
+          v9:BasicObject = LoadArg :arr@1
+          v10:BasicObject = LoadArg :index@2
+          v11:BasicObject = LoadArg :len@3
+          v12:BasicObject = LoadArg :val@4
+          Jump bb3(v8, v9, v10, v11, v12)
+        bb3(v14:BasicObject, v15:BasicObject, v16:BasicObject, v17:BasicObject, v18:BasicObject):
+          PatchPoint NoSingletonClass(Array@0x1000)
+          PatchPoint MethodRedefined(Array@0x1000, []=@0x1008, cme:0x1010)
+          v39:ArrayExact = GuardType v15, ArrayExact
+          v40:Fixnum = GuardType v16, Fixnum
+          v41:Fixnum = GuardType v17, Fixnum
+          v42:CUInt64 = LoadField v39, :_rbasic_flags@0x1038
+          v43:CUInt64 = GuardNoBitsSet v42, RUBY_FL_FREEZE=CUInt64(2048)
+          v44:CUInt64 = LoadField v39, :_rbasic_flags@0x1038
+          v45:CUInt64 = GuardNoBitsSet v44, RUBY_ELTS_SHARED=CUInt64(4096)
+          v46:CInt64 = UnboxFixnum v40
+          v47:CInt64 = UnboxFixnum v41
+          v48:BasicObject = CCall v39, :[]=@0x1040, v46, v47, v18
+          IncrCounter inline_cfunc_optimized_send_count
+          CheckInterrupts
+          Return v18
+        ");
+    }
+
+    #[test]
+    fn test_optimize_array_splice_aset_array_subclass() {
+        eval("
+            class MySpliceArray2 < Array; end
+            def test(arr, index, len, val)
+              arr[index, len] = val
+            end
+            a = MySpliceArray2.new([1,2,3])
+            test(a, 0, 2, [4,5])
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:4:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:BasicObject = GetLocal :arr, l0, SP@7
+          v3:BasicObject = GetLocal :index, l0, SP@6
+          v4:BasicObject = GetLocal :len, l0, SP@5
+          v5:BasicObject = GetLocal :val, l0, SP@4
+          Jump bb3(v1, v2, v3, v4, v5)
+        bb2():
+          EntryPoint JIT(0)
+          v8:BasicObject = LoadArg :self@0
+          v9:BasicObject = LoadArg :arr@1
+          v10:BasicObject = LoadArg :index@2
+          v11:BasicObject = LoadArg :len@3
+          v12:BasicObject = LoadArg :val@4
+          Jump bb3(v8, v9, v10, v11, v12)
+        bb3(v14:BasicObject, v15:BasicObject, v16:BasicObject, v17:BasicObject, v18:BasicObject):
+          PatchPoint NoSingletonClass(MySpliceArray2@0x1000)
+          PatchPoint MethodRedefined(MySpliceArray2@0x1000, []=@0x1008, cme:0x1010)
+          v39:ArraySubclass[class_exact:MySpliceArray2] = GuardType v15, ArraySubclass[class_exact:MySpliceArray2]
+          v40:BasicObject = CCallVariadic v39, :Array#[]=@0x1038, v16, v17, v18
+          CheckInterrupts
+          Return v18
+        ");
+    }
+
     #[test]
     fn test_optimize_array_ltlt() {
         eval("