MatchBranchSimplification: avoid intermediate vec allocation

This commit is contained in:
Tomasz Miąsko 2020-08-15 00:00:00 +00:00
parent 1fe2e2941d
commit af9b9e4ec8
2 changed files with 46 additions and 38 deletions

View file

@ -669,6 +669,17 @@ impl<I: Idx, T> IndexVec<I, T> {
}
}
/// Returns mutable references to three distinct elements or panics otherwise.
#[inline]
pub fn pick3_mut(&mut self, a: I, b: I, c: I) -> (&mut T, &mut T, &mut T) {
let (ai, bi, ci) = (a.index(), b.index(), c.index());
assert!(ai != bi && bi != ci && ci != ai);
let len = self.raw.len();
assert!(ai < len && bi < len && ci < len);
let ptr = self.raw.as_mut_ptr();
unsafe { (&mut *ptr.add(ai), &mut *ptr.add(bi), &mut *ptr.add(ci)) }
}
pub fn convert_index_type<Ix: Idx>(self) -> IndexVec<Ix, T> {
IndexVec { raw: self.raw, _marker: PhantomData }
}

View file

@ -48,7 +48,7 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
ref targets,
ref values,
..
} if targets.len() == 2 && values.len() == 1 => {
} if targets.len() == 2 && values.len() == 1 && targets[0] != targets[1] => {
(place, values[0], switch_ty, targets[0], targets[1])
}
// Only optimize switch int statements
@ -89,48 +89,45 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
// Take ownership of items now that we know we can optimize.
let discr = discr.clone();
let new_stmts = first_stmts
.iter()
.zip(scnd_stmts.iter())
.map(|(f, s)| {
match (&f.kind, &s.kind) {
(f_s, s_s) if f_s == s_s => (*f).clone(),
// We already checked that first and second are different blocks,
// and bb_idx has a different terminator from both of them.
let (from, first, second) = bbs.pick3_mut(bb_idx, first, second);
(
StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))),
) => {
// From earlier loop we know that we are dealing with bool constants only:
let f_b = f_c.literal.try_eval_bool(tcx, param_env).unwrap();
let s_b = s_c.literal.try_eval_bool(tcx, param_env).unwrap();
if f_b == s_b {
// Same value in both blocks. Use statement as is.
(*f).clone()
} else {
// Different value between blocks. Make value conditional on switch condition.
let size = tcx.layout_of(param_env.and(switch_ty)).unwrap().size;
let const_cmp = Operand::const_from_scalar(
tcx,
switch_ty,
crate::interpret::Scalar::from_uint(val, size),
rustc_span::DUMMY_SP,
);
let op = if f_b { BinOp::Eq } else { BinOp::Ne };
let rhs =
Rvalue::BinaryOp(op, Operand::Copy(discr.clone()), const_cmp);
Statement {
source_info: f.source_info,
kind: StatementKind::Assign(box (*lhs, rhs)),
}
let new_stmts = first.statements.iter().zip(second.statements.iter()).map(|(f, s)| {
match (&f.kind, &s.kind) {
(f_s, s_s) if f_s == s_s => (*f).clone(),
(
StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))),
) => {
// From earlier loop we know that we are dealing with bool constants only:
let f_b = f_c.literal.try_eval_bool(tcx, param_env).unwrap();
let s_b = s_c.literal.try_eval_bool(tcx, param_env).unwrap();
if f_b == s_b {
// Same value in both blocks. Use statement as is.
(*f).clone()
} else {
// Different value between blocks. Make value conditional on switch condition.
let size = tcx.layout_of(param_env.and(switch_ty)).unwrap().size;
let const_cmp = Operand::const_from_scalar(
tcx,
switch_ty,
crate::interpret::Scalar::from_uint(val, size),
rustc_span::DUMMY_SP,
);
let op = if f_b { BinOp::Eq } else { BinOp::Ne };
let rhs = Rvalue::BinaryOp(op, Operand::Copy(discr.clone()), const_cmp);
Statement {
source_info: f.source_info,
kind: StatementKind::Assign(box (*lhs, rhs)),
}
}
_ => unreachable!(),
}
})
.collect::<Vec<_>>();
let (from, first) = bbs.pick2_mut(bb_idx, first);
_ => unreachable!(),
}
});
from.statements.extend(new_stmts);
from.terminator_mut().kind = first.terminator().kind.clone();
}