Skip to content

Commit 96819cf

Browse files
tekknolagitekknolagi
committed
ZJIT: Parse splatarray, concattoarray, pushtoarray into HIR
1 parent 1435ea7 commit 96819cf

File tree

1 file changed

+148
-5
lines changed

1 file changed

+148
-5
lines changed

‎zjit/src/hir.rs‎

Lines changed: 148 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -326,12 +326,21 @@ pub enum Insn {
326326
StringCopy { val: InsnId },
327327
StringIntern { val: InsnId },
328328

329+
/// Call `to_a` on `val` if the method is defined, or make a new array `[val]` otherwise.
330+
ToArray { val: InsnId, state: InsnId },
331+
/// Call `to_a` on `val` if the method is defined, or make a new array `[val]` otherwise. If we
332+
/// called `to_a`, duplicate the returned array.
333+
ToNewArray { val: InsnId, state: InsnId },
329334
NewArray { elements: Vec<InsnId>, state: InsnId },
330335
/// NewHash contains a vec of (key, value) pairs
331336
NewHash { elements: Vec<(InsnId,InsnId)>, state: InsnId },
332337
ArraySet { array: InsnId, idx: usize, val: InsnId },
333338
ArrayDup { val: InsnId, state: InsnId },
334339
ArrayMax { elements: Vec<InsnId>, state: InsnId },
340+
/// Extend `left` with the elements from `right`. `left` and `right` must both be `Array`.
341+
ArrayExtend { left: InsnId, right: InsnId, state: InsnId },
342+
/// Push `val` onto `array`, where `array` is already `Array`.
343+
ArrayPush { array: InsnId, val: InsnId, state: InsnId },
335344

336345
HashDup { val: InsnId, state: InsnId },
337346

@@ -401,7 +410,8 @@ impl Insn {
401410
match self {
402411
Insn::ArraySet { .. } | Insn::Snapshot { .. } | Insn::Jump(_)
403412
| Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. }
404-
| Insn::PatchPoint { .. } | Insn::SetIvar { .. } => false,
413+
| Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
414+
| Insn::ArrayPush { .. } => false,
405415
_ => true,
406416
}
407417
}
@@ -546,6 +556,10 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
546556
Insn::Snapshot { state } => write!(f, "Snapshot {}", state),
547557
Insn::GetIvar { self_val, id, .. } => write!(f, "GetIvar {self_val}, :{}", id.contents_lossy().into_owned()),
548558
Insn::SetIvar { self_val, id, val, .. } => write!(f, "SetIvar {self_val}, :{}, {val}", id.contents_lossy().into_owned()),
559+
Insn::ToArray { val, .. } => write!(f, "ToArray {val}"),
560+
Insn::ToNewArray { val, .. } => write!(f, "ToNewArray {val}"),
561+
Insn::ArrayExtend { left, right, .. } => write!(f, "ArrayExtend {left}, {right}"),
562+
Insn::ArrayPush { array, val, .. } => write!(f, "ArrayPush {array}, {val}"),
549563
insn => { write!(f, "{insn:?}") }
550564
}
551565
}
@@ -897,6 +911,10 @@ impl Function {
897911
ArrayMax { elements, state } => ArrayMax { elements: find_vec!(*elements), state: find!(*state) },
898912
&GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
899913
&SetIvar { self_val, id, val, state } => SetIvar { self_val: find!(self_val), id, val, state },
914+
&ToArray { val, state } => ToArray { val: find!(val), state },
915+
&ToNewArray { val, state } => ToNewArray { val: find!(val), state },
916+
&ArrayExtend { left, right, state } => ArrayExtend { left: find!(left), right: find!(right), state },
917+
&ArrayPush { array, val, state } => ArrayPush { array: find!(array), val: find!(val), state },
900918
}
901919
}
902920

@@ -922,7 +940,8 @@ impl Function {
922940
Insn::Param { .. } => unimplemented!("params should not be present in block.insns"),
923941
Insn::ArraySet { .. } | Insn::Snapshot { .. } | Insn::Jump(_)
924942
| Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. }
925-
| Insn::PatchPoint { .. } | Insn::SetIvar { .. } =>
943+
| Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
944+
| Insn::ArrayPush { .. } =>
926945
panic!("Cannot infer type of instruction with no output"),
927946
Insn::Const { val: Const::Value(val) } => Type::from_value(*val),
928947
Insn::Const { val: Const::CBool(val) } => Type::from_cbool(*val),
@@ -967,6 +986,8 @@ impl Function {
967986
Insn::GetConstantPath { .. } => types::BasicObject,
968987
Insn::ArrayMax { .. } => types::BasicObject,
969988
Insn::GetIvar { .. } => types::BasicObject,
989+
Insn::ToNewArray { .. } => types::ArrayExact,
990+
Insn::ToArray { .. } => types::ArrayExact,
970991
}
971992
}
972993

@@ -1431,7 +1452,9 @@ impl Function {
14311452
| Insn::Test { val } =>
14321453
worklist.push_back(val),
14331454
Insn::GuardType { val, state, .. }
1434-
| Insn::GuardBitEquals { val, state, .. } => {
1455+
| Insn::GuardBitEquals { val, state, .. }
1456+
| Insn::ToArray { val, state }
1457+
| Insn::ToNewArray { val, state } => {
14351458
worklist.push_back(val);
14361459
worklist.push_back(state);
14371460
}
@@ -1448,6 +1471,7 @@ impl Function {
14481471
| Insn::FixnumMult { left, right, state }
14491472
| Insn::FixnumDiv { left, right, state }
14501473
| Insn::FixnumMod { left, right, state }
1474+
| Insn::ArrayExtend { left, right, state }
14511475
=> {
14521476
worklist.push_back(left);
14531477
worklist.push_back(right);
@@ -1489,6 +1513,11 @@ impl Function {
14891513
worklist.push_back(val);
14901514
worklist.push_back(state);
14911515
}
1516+
Insn::ArrayPush { array, val, state } => {
1517+
worklist.push_back(array);
1518+
worklist.push_back(val);
1519+
worklist.push_back(state);
1520+
}
14921521
}
14931522
}
14941523
// Now remove all unnecessary instructions
@@ -1976,6 +2005,39 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
19762005
let insn_id = fun.push_insn(block, Insn::HashDup { val, state: exit_id });
19772006
state.stack_push(insn_id);
19782007
}
2008+
YARVINSN_splatarray => {
2009+
let flag = get_arg(pc, 0);
2010+
let result_must_be_mutable = flag.test();
2011+
let val = state.stack_pop()?;
2012+
let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
2013+
let obj = if result_must_be_mutable {
2014+
fun.push_insn(block, Insn::ToNewArray { val, state: exit_id })
2015+
} else {
2016+
fun.push_insn(block, Insn::ToArray { val, state: exit_id })
2017+
};
2018+
state.stack_push(obj);
2019+
}
2020+
YARVINSN_concattoarray => {
2021+
let right = state.stack_pop()?;
2022+
let left = state.stack_pop()?;
2023+
let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
2024+
let right_array = fun.push_insn(block, Insn::ToArray { val: right, state: exit_id });
2025+
fun.push_insn(block, Insn::ArrayExtend { left, right: right_array, state: exit_id });
2026+
state.stack_push(left);
2027+
}
2028+
YARVINSN_pushtoarray => {
2029+
let count = get_arg(pc, 0).as_usize();
2030+
let mut vals = vec![];
2031+
for _ in 0..count {
2032+
vals.push(state.stack_pop()?);
2033+
}
2034+
let array = state.stack_pop()?;
2035+
let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
2036+
for val in vals.into_iter().rev() {
2037+
fun.push_insn(block, Insn::ArrayPush { array, val, state: exit_id });
2038+
}
2039+
state.stack_push(array);
2040+
}
19792041
YARVINSN_putobject_INT2FIX_0_ => {
19802042
state.stack_push(fun.push_insn(block, Insn::Const { val: Const::Value(VALUE::fixnum_from_usize(0)) }));
19812043
}
@@ -3006,7 +3068,7 @@ mod tests {
30063068
eval("
30073069
def test(a) = foo(*a)
30083070
");
3009-
assert_compile_fails("test", ParseError::UnknownOpcode("splatarray".into()))
3071+
assert_compile_fails("test", ParseError::UnhandledCallType(CallType::Splat))
30103072
}
30113073

30123074
#[test]
@@ -3074,7 +3136,7 @@ mod tests {
30743136
eval("
30753137
def test(*) = foo *, 1
30763138
");
3077-
assert_compile_fails("test", ParseError::UnknownOpcode("splatarray".into()))
3139+
assert_compile_fails("test", ParseError::UnhandledCallType(CallType::SplatMut))
30783140
}
30793141

30803142
#[test]
@@ -3192,6 +3254,69 @@ mod tests {
31923254
Return v1
31933255
"#]]);
31943256
}
3257+
3258+
#[test]
3259+
fn test_splatarray_mut() {
3260+
eval("
3261+
def test(a) = [*a]
3262+
");
3263+
assert_method_hir("test", expect![[r#"
3264+
fn test:
3265+
bb0(v0:BasicObject):
3266+
v3:ArrayExact = ToNewArray v0
3267+
Return v3
3268+
"#]]);
3269+
}
3270+
3271+
#[test]
3272+
fn test_concattoarray() {
3273+
eval("
3274+
def test(a) = [1, *a]
3275+
");
3276+
assert_method_hir("test", expect![[r#"
3277+
fn test:
3278+
bb0(v0:BasicObject):
3279+
v2:Fixnum[1] = Const Value(1)
3280+
v4:ArrayExact = NewArray v2
3281+
v6:ArrayExact = ToArray v0
3282+
ArrayExtend v4, v6
3283+
Return v4
3284+
"#]]);
3285+
}
3286+
3287+
#[test]
3288+
fn test_pushtoarray_one_element() {
3289+
eval("
3290+
def test(a) = [*a, 1]
3291+
");
3292+
assert_method_hir("test", expect![[r#"
3293+
fn test:
3294+
bb0(v0:BasicObject):
3295+
v3:ArrayExact = ToNewArray v0
3296+
v4:Fixnum[1] = Const Value(1)
3297+
ArrayPush v3, v4
3298+
Return v3
3299+
"#]]);
3300+
}
3301+
3302+
#[test]
3303+
fn test_pushtoarray_multiple_elements() {
3304+
eval("
3305+
def test(a) = [*a, 1, 2, 3]
3306+
");
3307+
assert_method_hir("test", expect![[r#"
3308+
fn test:
3309+
bb0(v0:BasicObject):
3310+
v3:ArrayExact = ToNewArray v0
3311+
v4:Fixnum[1] = Const Value(1)
3312+
v5:Fixnum[2] = Const Value(2)
3313+
v6:Fixnum[3] = Const Value(3)
3314+
ArrayPush v3, v4
3315+
ArrayPush v3, v5
3316+
ArrayPush v3, v6
3317+
Return v3
3318+
"#]]);
3319+
}
31953320
}
31963321

31973322
#[cfg(test)]
@@ -3719,6 +3844,24 @@ mod opt_tests {
37193844
"#]]);
37203845
}
37213846

3847+
#[test]
3848+
fn test_do_not_eliminate_to_new_array() {
3849+
eval("
3850+
def test(a)
3851+
c = [*a]
3852+
5
3853+
end
3854+
");
3855+
assert_optimized_method_hir("test", expect![[r#"
3856+
fn test:
3857+
bb0(v0:BasicObject):
3858+
v1:NilClassExact = Const Value(nil)
3859+
v4:ArrayExact = ToNewArray v0
3860+
v5:Fixnum[5] = Const Value(5)
3861+
Return v5
3862+
"#]]);
3863+
}
3864+
37223865
#[test]
37233866
fn test_eliminate_hash_dup() {
37243867
eval("

0 commit comments

Comments
 (0)