Skip to content

Commit f73ca40

Browse files
nmeansnmeans
committed
ZJIT: Use dedicated ArraySplice HIR instruction for splice safety
Replace bare CCall to rb_jit_ary_aset_by_rb_ary_splice with a dedicated ArraySplice instruction that carries frame state, following the HashAset/ArrayPop pattern. Codegen uses gen_prepare_non_leaf_call to save PC/SP and spill locals before the C call, ensuring correct GC and exception handling. Also add tests for IndexError (negative length, out-of-bounds index) and single-value splice (exercises rb_ary_to_ary allocation path).
1 parent ea4d991 commit f73ca40

File tree

5 files changed

+77
-10
lines changed

5 files changed

+77
-10
lines changed

‎zjit/src/codegen.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -480,6 +480,9 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
480480
Insn::ArrayAset { array, index, val } => {
481481
no_output!(gen_array_aset(asm, opnd!(array), opnd!(index), opnd!(val)))
482482
}
483+
Insn::ArraySplice { array, beg, len, val, state } => {
484+
no_output!(gen_array_splice(jit, asm, opnd!(array), opnd!(beg), opnd!(len), opnd!(val), &function.frame_state(*state)))
485+
}
483486
Insn::ArrayPop { array, state } => gen_array_pop(asm, opnd!(array), &function.frame_state(*state)),
484487
Insn::ArrayLength { array } => gen_array_length(asm, opnd!(array)),
485488
Insn::ObjectAlloc { val, state } => gen_object_alloc(jit, asm, opnd!(val), &function.frame_state(*state)),
@@ -1712,6 +1715,11 @@ fn gen_array_aset(
17121715
asm.store(Opnd::mem(VALUE_BITS, elem_ptr, 0), val);
17131716
}
17141717

1718+
fn gen_array_splice(jit: &JITState, asm: &mut Assembler, array: Opnd, beg: Opnd, len: Opnd, val: Opnd, state: &FrameState) {
1719+
gen_prepare_non_leaf_call(jit, asm, state);
1720+
asm_ccall!(asm, rb_jit_ary_aset_by_rb_ary_splice, array, beg, len, val);
1721+
}
1722+
17151723
fn gen_array_pop(asm: &mut Assembler, array: Opnd, state: &FrameState) -> lir::Opnd {
17161724
gen_prepare_leaf_call_with_gc(asm, state);
17171725
asm_ccall!(asm, rb_ary_pop, array)

‎zjit/src/codegen_tests.rs‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2847,6 +2847,49 @@ fn test_array_splice_aset_array_subclass() {
28472847
assert_snapshot!(inspect("arr = MySpliceArray.new([1,2,3]); test(arr, 0, 2, [4,5]); arr.to_a"), @"[4, 5, 3]");
28482848
}
28492849

2850+
#[test]
2851+
fn test_array_splice_aset_negative_length() {
2852+
assert_snapshot!(inspect("
2853+
def test(arr, idx, len, val)
2854+
arr[idx, len] = val
2855+
end
2856+
arr = [1,2,3]
2857+
test(arr, 0, 2, [4,5])
2858+
begin
2859+
test(arr, 0, -1, [4,5])
2860+
rescue => e
2861+
e.class
2862+
end
2863+
"), @"IndexError");
2864+
}
2865+
2866+
#[test]
2867+
fn test_array_splice_aset_index_out_of_bounds() {
2868+
assert_snapshot!(inspect("
2869+
def test(arr, idx, len, val)
2870+
arr[idx, len] = val
2871+
end
2872+
arr = [1,2,3]
2873+
test(arr, 0, 2, [4,5])
2874+
begin
2875+
test(arr, -100, 2, [4,5])
2876+
rescue => e
2877+
e.class
2878+
end
2879+
"), @"IndexError");
2880+
}
2881+
2882+
#[test]
2883+
fn test_array_splice_aset_single_value() {
2884+
eval("
2885+
def test(arr, idx, len, val)
2886+
arr[idx, len] = val
2887+
end
2888+
test([1,2,3], 0, 2, 42)
2889+
");
2890+
assert_snapshot!(inspect("arr = [1,2,3]; test(arr, 0, 2, 42); arr"), @"[42, 3]");
2891+
}
2892+
28502893
#[test]
28512894
fn test_array_aset_non_fixnum_index() {
28522895
assert_snapshot!(inspect(r#"

‎zjit/src/cruby_methods.rs‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -392,13 +392,8 @@ fn inline_array_aset(fun: &mut hir::Function, block: hir::BlockId, recv: hir::In
392392

393393
let index = fun.push_insn(block, hir::Insn::UnboxFixnum { val: index });
394394
let len = fun.push_insn(block, hir::Insn::UnboxFixnum { val: len });
395-
let _ = fun.push_insn(block, hir::Insn::CCall {
396-
cfunc: rb_jit_ary_aset_by_rb_ary_splice as _,
397-
recv,
398-
args: vec![index, len, val],
399-
return_type: types::BasicObject,
400-
elidable: false,
401-
name: ID!(aset),
395+
let _ = fun.push_insn(block, hir::Insn::ArraySplice {
396+
array: recv, beg: index, len, val, state,
402397
});
403398
return Some(val);
404399
}

‎zjit/src/hir.rs‎

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -777,6 +777,8 @@ pub enum Insn {
777777
ArrayPush { array: InsnId, val: InsnId, state: InsnId },
778778
ArrayAref { array: InsnId, index: InsnId },
779779
ArrayAset { array: InsnId, index: InsnId, val: InsnId },
780+
/// Splice val into array at [beg, len]. Can allocate and raise IndexError.
781+
ArraySplice { array: InsnId, beg: InsnId, len: InsnId, val: InsnId, state: InsnId },
780782
ArrayPop { array: InsnId, state: InsnId },
781783
/// Return the length of the array as a C `long` ([`types::CInt64`])
782784
ArrayLength { array: InsnId },
@@ -1070,7 +1072,7 @@ impl Insn {
10701072
| Insn::SetLocal { .. } | Insn::Throw { .. } | Insn::IncrCounter(_) | Insn::IncrCounterPtr { .. }
10711073
| Insn::CheckInterrupts { .. }
10721074
| Insn::StoreField { .. } | Insn::WriteBarrier { .. } | Insn::HashAset { .. }
1073-
| Insn::ArrayAset { .. } => false,
1075+
| Insn::ArrayAset { .. } | Insn::ArraySplice { .. } => false,
10741076
_ => true,
10751077
}
10761078
}
@@ -1137,6 +1139,7 @@ impl Insn {
11371139
Insn::ArrayPush { .. } => effects::Any,
11381140
Insn::ArrayAref { .. } => effects::Any,
11391141
Insn::ArrayAset { .. } => effects::Any,
1142+
Insn::ArraySplice { .. } => effects::Any,
11401143
Insn::ArrayPop { .. } => effects::Any,
11411144
Insn::ArrayLength { .. } => Effect::write(abstract_heaps::Empty),
11421145
Insn::HashAref { .. } => effects::Any,
@@ -1347,6 +1350,9 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
13471350
Insn::ArrayAset { array, index, val, ..} => {
13481351
write!(f, "ArrayAset {array}, {index}, {val}")
13491352
}
1353+
Insn::ArraySplice { array, beg, len, val, .. } => {
1354+
write!(f, "ArraySplice {array}, {beg}, {len}, {val}")
1355+
}
13501356
Insn::ArrayPop { array, .. } => {
13511357
write!(f, "ArrayPop {array}")
13521358
}
@@ -2386,6 +2392,7 @@ impl Function {
23862392
&NewRangeFixnum { low, high, flag, state } => NewRangeFixnum { low: find!(low), high: find!(high), flag, state: find!(state) },
23872393
&ArrayAref { array, index } => ArrayAref { array: find!(array), index: find!(index) },
23882394
&ArrayAset { array, index, val } => ArrayAset { array: find!(array), index: find!(index), val: find!(val) },
2395+
&ArraySplice { array, beg, len, val, state } => ArraySplice { array: find!(array), beg: find!(beg), len: find!(len), val: find!(val), state },
23892396
&ArrayPop { array, state } => ArrayPop { array: find!(array), state: find!(state) },
23902397
&ArrayLength { array } => ArrayLength { array: find!(array) },
23912398
&ArrayMax { ref elements, state } => ArrayMax { elements: find_vec!(elements), state: find!(state) },
@@ -2460,7 +2467,8 @@ impl Function {
24602467
| Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. }
24612468
| Insn::IncrCounter(_) | Insn::IncrCounterPtr { .. }
24622469
| Insn::CheckInterrupts { .. }
2463-
| Insn::StoreField { .. } | Insn::WriteBarrier { .. } | Insn::HashAset { .. } | Insn::ArrayAset { .. } =>
2470+
| Insn::StoreField { .. } | Insn::WriteBarrier { .. } | Insn::HashAset { .. } | Insn::ArrayAset { .. }
2471+
| Insn::ArraySplice { .. } =>
24642472
panic!("Cannot infer type of instruction with no output: {}. See Insn::has_output().", self.insns[insn.0]),
24652473
Insn::Const { val: Const::Value(val) } => Type::from_value(*val),
24662474
Insn::Const { val: Const::CBool(val) } => Type::from_cbool(*val),
@@ -4802,6 +4810,13 @@ impl Function {
48024810
worklist.push_back(index);
48034811
worklist.push_back(val);
48044812
}
4813+
&Insn::ArraySplice { array, beg, len, val, state } => {
4814+
worklist.push_back(array);
4815+
worklist.push_back(beg);
4816+
worklist.push_back(len);
4817+
worklist.push_back(val);
4818+
worklist.push_back(state);
4819+
}
48054820
&Insn::ArrayPop { array, state } => {
48064821
worklist.push_back(array);
48074822
worklist.push_back(state);
@@ -5619,6 +5634,12 @@ impl Function {
56195634
self.assert_subtype(insn_id, array, types::ArrayExact)?;
56205635
self.assert_subtype(insn_id, index, types::CInt64)
56215636
}
5637+
Insn::ArraySplice { array, beg, len, val, .. } => {
5638+
self.assert_subtype(insn_id, array, types::ArrayExact)?;
5639+
self.assert_subtype(insn_id, beg, types::CInt64)?;
5640+
self.assert_subtype(insn_id, len, types::CInt64)?;
5641+
self.assert_subtype(insn_id, val, types::BasicObject)
5642+
}
56225643
// Instructions with Hash operands
56235644
Insn::HashAref { hash, .. }
56245645
| Insn::HashAset { hash, .. } => self.assert_subtype(insn_id, hash, types::HashExact),

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8421,7 +8421,7 @@ mod hir_opt_tests {
84218421
v45:CUInt64 = GuardNoBitsSet v44, RUBY_ELTS_SHARED=CUInt64(4096)
84228422
v46:CInt64 = UnboxFixnum v40
84238423
v47:CInt64 = UnboxFixnum v41
8424-
v48:BasicObject = CCall v39, :[]=@0x1040, v46, v47, v18
8424+
ArraySplice v39, v46, v47, v18
84258425
IncrCounter inline_cfunc_optimized_send_count
84268426
CheckInterrupts
84278427
Return v18

0 commit comments

Comments
 (0)