Skip to content

Commit 309d6ef

Browse files
st0012st0012
committed
ZJIT: Inline Hash#[]= 
1 parent 3a0596b commit 309d6ef

File tree

5 files changed

+194
-3
lines changed

5 files changed

+194
-3
lines changed

‎test/ruby/test_zjit.rb‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1429,6 +1429,77 @@ def test = {}.freeze
14291429
}, insns: [:opt_hash_freeze], call_threshold: 1
14301430
end
14311431

1432+
def test_opt_aset_hash
1433+
assert_compiles '42', %q{
1434+
def test(h, k, v)
1435+
h[k] = v
1436+
end
1437+
h = {}
1438+
test(h, :key, 42)
1439+
test(h, :key, 42)
1440+
h[:key]
1441+
}, call_threshold: 2, insns: [:opt_aset]
1442+
end
1443+
1444+
def test_opt_aset_hash_returns_value
1445+
assert_compiles '100', %q{
1446+
def test(h, k, v)
1447+
h[k] = v
1448+
end
1449+
test({}, :key, 100)
1450+
test({}, :key, 100)
1451+
}, call_threshold: 2
1452+
end
1453+
1454+
def test_opt_aset_hash_string_key
1455+
assert_compiles '"bar"', %q{
1456+
def test(h, k, v)
1457+
h[k] = v
1458+
end
1459+
h = {}
1460+
test(h, "foo", "bar")
1461+
test(h, "foo", "bar")
1462+
h["foo"]
1463+
}, call_threshold: 2
1464+
end
1465+
1466+
def test_opt_aset_hash_subclass
1467+
assert_compiles '42', %q{
1468+
class MyHash < Hash; end
1469+
def test(h, k, v)
1470+
h[k] = v
1471+
end
1472+
h = MyHash.new
1473+
test(h, :key, 42)
1474+
test(h, :key, 42)
1475+
h[:key]
1476+
}, call_threshold: 2
1477+
end
1478+
1479+
def test_opt_aset_hash_too_few_args
1480+
assert_compiles '"ArgumentError"', %q{
1481+
def test(h)
1482+
h.[]= 123
1483+
rescue ArgumentError
1484+
"ArgumentError"
1485+
end
1486+
test({})
1487+
test({})
1488+
}, call_threshold: 2
1489+
end
1490+
1491+
def test_opt_aset_hash_too_many_args
1492+
assert_compiles '"ArgumentError"', %q{
1493+
def test(h)
1494+
h[:a, :b] = :c
1495+
rescue ArgumentError
1496+
"ArgumentError"
1497+
end
1498+
test({})
1499+
test({})
1500+
}, call_threshold: 2
1501+
end
1502+
14321503
def test_opt_ary_freeze
14331504
assert_compiles "[[], 5]", %q{
14341505
def test = [].freeze

‎zjit/src/codegen.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -484,6 +484,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
484484
&Insn::CheckInterrupts { state } => no_output!(gen_check_interrupts(jit, asm, &function.frame_state(state))),
485485
&Insn::HashDup { val, state } => { gen_hash_dup(asm, opnd!(val), &function.frame_state(state)) },
486486
&Insn::HashAref { hash, key, state } => { gen_hash_aref(jit, asm, opnd!(hash), opnd!(key), &function.frame_state(state)) },
487+
&Insn::HashAset { hash, key, val, state } => { no_output!(gen_hash_aset(jit, asm, opnd!(hash), opnd!(key), opnd!(val), &function.frame_state(state))) },
487488
&Insn::ArrayPush { array, val, state } => { no_output!(gen_array_push(asm, opnd!(array), opnd!(val), &function.frame_state(state))) },
488489
&Insn::ToNewArray { val, state } => { gen_to_new_array(jit, asm, opnd!(val), &function.frame_state(state)) },
489490
&Insn::ToArray { val, state } => { gen_to_array(jit, asm, opnd!(val), &function.frame_state(state)) },
@@ -1065,6 +1066,11 @@ fn gen_hash_aref(jit: &mut JITState, asm: &mut Assembler, hash: Opnd, key: Opnd,
10651066
asm_ccall!(asm, rb_hash_aref, hash, key)
10661067
}
10671068

1069+
fn gen_hash_aset(jit: &mut JITState, asm: &mut Assembler, hash: Opnd, key: Opnd, val: Opnd, state: &FrameState) {
1070+
gen_prepare_non_leaf_call(jit, asm, state);
1071+
asm_ccall!(asm, rb_hash_aset, hash, key, val);
1072+
}
1073+
10681074
fn gen_array_push(asm: &mut Assembler, array: Opnd, val: Opnd, state: &FrameState) {
10691075
gen_prepare_leaf_call_with_gc(asm, state);
10701076
asm_ccall!(asm, rb_ary_push, array, val);

‎zjit/src/cruby_methods.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@ pub fn init() -> Annotations {
229229
annotate!(rb_cArray, "push", inline_array_push);
230230
annotate!(rb_cArray, "pop", inline_array_pop);
231231
annotate!(rb_cHash, "[]", inline_hash_aref);
232+
annotate!(rb_cHash, "[]=", inline_hash_aset);
232233
annotate!(rb_cHash, "size", types::Fixnum, no_gc, leaf, elidable);
233234
annotate!(rb_cHash, "empty?", types::BoolExact, no_gc, leaf, elidable);
234235
annotate!(rb_cNilClass, "nil?", inline_nilclass_nil_p);
@@ -356,6 +357,12 @@ fn inline_hash_aref(fun: &mut hir::Function, block: hir::BlockId, recv: hir::Ins
356357
None
357358
}
358359

360+
fn inline_hash_aset(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
361+
let &[key, val] = args else { return None; };
362+
let _ = fun.push_insn(block, hir::Insn::HashAset { hash: recv, key, val, state });
363+
// Hash#[]= returns the value, not the hash
364+
Some(val)
365+
}
359366

360367
fn inline_string_bytesize(fun: &mut hir::Function, block: hir::BlockId, recv: hir::InsnId, args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
361368
if args.is_empty() && fun.likely_a(recv, types::String, state) {

‎zjit/src/hir.rs‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -733,6 +733,7 @@ pub enum Insn {
733733
ArrayLength { array: InsnId },
734734

735735
HashAref { hash: InsnId, key: InsnId, state: InsnId },
736+
HashAset { hash: InsnId, key: InsnId, val: InsnId, state: InsnId },
736737
HashDup { val: InsnId, state: InsnId },
737738

738739
/// Allocate an instance of the `val` object without calling `#initialize` on it.
@@ -991,7 +992,8 @@ impl Insn {
991992
| Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::SetClassVar { .. } | Insn::ArrayExtend { .. }
992993
| Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetGlobal { .. }
993994
| Insn::SetLocal { .. } | Insn::Throw { .. } | Insn::IncrCounter(_) | Insn::IncrCounterPtr { .. }
994-
| Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::StoreField { .. } | Insn::WriteBarrier { .. } => false,
995+
| Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::StoreField { .. } | Insn::WriteBarrier { .. }
996+
| Insn::HashAset { .. } => false,
995997
_ => true,
996998
}
997999
}
@@ -1182,6 +1184,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
11821184
Insn::ArrayDup { val, .. } => { write!(f, "ArrayDup {val}") }
11831185
Insn::HashDup { val, .. } => { write!(f, "HashDup {val}") }
11841186
Insn::HashAref { hash, key, .. } => { write!(f, "HashAref {hash}, {key}")}
1187+
Insn::HashAset { hash, key, val, .. } => { write!(f, "HashAset {hash}, {key}, {val}")}
11851188
Insn::ObjectAlloc { val, .. } => { write!(f, "ObjectAlloc {val}") }
11861189
&Insn::ObjectAllocClass { class, .. } => {
11871190
let class_name = get_class_name(class);
@@ -2034,6 +2037,7 @@ impl Function {
20342037
&ArrayDup { val, state } => ArrayDup { val: find!(val), state },
20352038
&HashDup { val, state } => HashDup { val: find!(val), state },
20362039
&HashAref { hash, key, state } => HashAref { hash: find!(hash), key: find!(key), state },
2040+
&HashAset { hash, key, val, state } => HashAset { hash: find!(hash), key: find!(key), val: find!(val), state },
20372041
&ObjectAlloc { val, state } => ObjectAlloc { val: find!(val), state },
20382042
&ObjectAllocClass { class, state } => ObjectAllocClass { class, state: find!(state) },
20392043
&CCall { cfunc, recv, ref args, name, return_type, elidable } => CCall { cfunc, recv: find!(recv), args: find_vec!(args), name, return_type, elidable },
@@ -2131,7 +2135,7 @@ impl Function {
21312135
| Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::SetClassVar { .. } | Insn::ArrayExtend { .. }
21322136
| Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetLocal { .. } | Insn::IncrCounter(_)
21332137
| Insn::CheckInterrupts { .. } | Insn::GuardBlockParamProxy { .. } | Insn::IncrCounterPtr { .. }
2134-
| Insn::StoreField { .. } | Insn::WriteBarrier { .. } =>
2138+
| Insn::StoreField { .. } | Insn::WriteBarrier { .. } | Insn::HashAset { .. } =>
21352139
panic!("Cannot infer type of instruction with no output: {}. See Insn::has_output().", self.insns[insn.0]),
21362140
Insn::Const { val: Const::Value(val) } => Type::from_value(*val),
21372141
Insn::Const { val: Const::CBool(val) } => Type::from_cbool(*val),
@@ -4004,6 +4008,12 @@ impl Function {
40044008
worklist.push_back(key);
40054009
worklist.push_back(state);
40064010
}
4011+
&Insn::HashAset { hash, key, val, state } => {
4012+
worklist.push_back(hash);
4013+
worklist.push_back(key);
4014+
worklist.push_back(val);
4015+
worklist.push_back(state);
4016+
}
40074017
&Insn::Send { recv, ref args, state, .. }
40084018
| &Insn::SendForward { recv, ref args, state, .. }
40094019
| &Insn::SendWithoutBlock { recv, ref args, state, .. }
@@ -4699,7 +4709,8 @@ impl Function {
46994709
self.assert_subtype(insn_id, index, types::Fixnum)
47004710
}
47014711
// Instructions with Hash operands
4702-
Insn::HashAref { hash, .. } => self.assert_subtype(insn_id, hash, types::Hash),
4712+
Insn::HashAref { hash, .. }
4713+
| Insn::HashAset { hash, .. } => self.assert_subtype(insn_id, hash, types::Hash),
47034714
Insn::HashDup { val, .. } => self.assert_subtype(insn_id, val, types::HashExact),
47044715
// Other
47054716
Insn::ObjectAllocClass { class, .. } => {

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6738,6 +6738,102 @@ mod hir_opt_tests {
67386738
");
67396739
}
67406740

6741+
#[test]
6742+
fn test_hash_aset_literal() {
6743+
eval("
6744+
def test
6745+
h = {}
6746+
h[1] = 3
6747+
end
6748+
");
6749+
assert_snapshot!(hir_string("test"), @r"
6750+
fn test@<compiled>:3:
6751+
bb0():
6752+
EntryPoint interpreter
6753+
v1:BasicObject = LoadSelf
6754+
v2:NilClass = Const Value(nil)
6755+
Jump bb2(v1, v2)
6756+
bb1(v5:BasicObject):
6757+
EntryPoint JIT(0)
6758+
v6:NilClass = Const Value(nil)
6759+
Jump bb2(v5, v6)
6760+
bb2(v8:BasicObject, v9:NilClass):
6761+
v13:HashExact = NewHash
6762+
PatchPoint NoEPEscape(test)
6763+
v22:Fixnum[1] = Const Value(1)
6764+
v24:Fixnum[3] = Const Value(3)
6765+
PatchPoint MethodRedefined(Hash@0x1000, []=@0x1008, cme:0x1010)
6766+
PatchPoint NoSingletonClass(Hash@0x1000)
6767+
HashAset v13, v22, v24
6768+
IncrCounter inline_cfunc_optimized_send_count
6769+
CheckInterrupts
6770+
Return v24
6771+
");
6772+
}
6773+
6774+
#[test]
6775+
fn test_hash_aset_profiled() {
6776+
eval("
6777+
def test(hash, key, val)
6778+
hash[key] = val
6779+
end
6780+
test({}, 0, 1)
6781+
");
6782+
assert_snapshot!(hir_string("test"), @r"
6783+
fn test@<compiled>:3:
6784+
bb0():
6785+
EntryPoint interpreter
6786+
v1:BasicObject = LoadSelf
6787+
v2:BasicObject = GetLocal :hash, l0, SP@6
6788+
v3:BasicObject = GetLocal :key, l0, SP@5
6789+
v4:BasicObject = GetLocal :val, l0, SP@4
6790+
Jump bb2(v1, v2, v3, v4)
6791+
bb1(v7:BasicObject, v8:BasicObject, v9:BasicObject, v10:BasicObject):
6792+
EntryPoint JIT(0)
6793+
Jump bb2(v7, v8, v9, v10)
6794+
bb2(v12:BasicObject, v13:BasicObject, v14:BasicObject, v15:BasicObject):
6795+
PatchPoint MethodRedefined(Hash@0x1000, []=@0x1008, cme:0x1010)
6796+
PatchPoint NoSingletonClass(Hash@0x1000)
6797+
v35:HashExact = GuardType v13, HashExact
6798+
HashAset v35, v14, v15
6799+
IncrCounter inline_cfunc_optimized_send_count
6800+
CheckInterrupts
6801+
Return v15
6802+
");
6803+
}
6804+
6805+
#[test]
6806+
fn test_hash_aset_subclass() {
6807+
eval("
6808+
class C < Hash; end
6809+
def test(hash, key, val)
6810+
hash[key] = val
6811+
end
6812+
test(C.new, 0, 1)
6813+
");
6814+
assert_snapshot!(hir_string("test"), @r"
6815+
fn test@<compiled>:4:
6816+
bb0():
6817+
EntryPoint interpreter
6818+
v1:BasicObject = LoadSelf
6819+
v2:BasicObject = GetLocal :hash, l0, SP@6
6820+
v3:BasicObject = GetLocal :key, l0, SP@5
6821+
v4:BasicObject = GetLocal :val, l0, SP@4
6822+
Jump bb2(v1, v2, v3, v4)
6823+
bb1(v7:BasicObject, v8:BasicObject, v9:BasicObject, v10:BasicObject):
6824+
EntryPoint JIT(0)
6825+
Jump bb2(v7, v8, v9, v10)
6826+
bb2(v12:BasicObject, v13:BasicObject, v14:BasicObject, v15:BasicObject):
6827+
PatchPoint MethodRedefined(C@0x1000, []=@0x1008, cme:0x1010)
6828+
PatchPoint NoSingletonClass(C@0x1000)
6829+
v35:HashSubclass[class_exact:C] = GuardType v13, HashSubclass[class_exact:C]
6830+
HashAset v35, v14, v15
6831+
IncrCounter inline_cfunc_optimized_send_count
6832+
CheckInterrupts
6833+
Return v15
6834+
");
6835+
}
6836+
67416837
#[test]
67426838
fn test_optimize_thread_current() {
67436839
eval("

0 commit comments

Comments
 (0)