Refactor how SwitchInt stores jump targets

This commit is contained in:
Jonas Schievink
2020-10-10 17:36:04 +02:00
parent cae8bc1f23
commit 432535da2b
22 changed files with 247 additions and 206 deletions

View File

@@ -1,7 +1,7 @@
use crate::{transform::MirPass, util::patch::MirPatch};
use rustc_middle::mir::*;
use rustc_middle::ty::{Ty, TyCtxt};
use std::{borrow::Cow, fmt::Debug};
use std::fmt::Debug;
use super::simplify::simplify_cfg;
@@ -95,15 +95,17 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
StatementKind::Assign(box (Place::from(not_equal_temp), not_equal_rvalue)),
);
let (mut targets_to_jump_to, values_to_jump_to): (Vec<_>, Vec<_>) = opt_to_apply
let new_targets = opt_to_apply
.infos
.iter()
.flat_map(|x| x.second_switch_info.targets_with_values.iter())
.cloned()
.unzip();
.cloned();
let targets = SwitchTargets::new(
new_targets,
opt_to_apply.infos[0].first_switch_info.otherwise_bb,
);
// add otherwise case in the end
targets_to_jump_to.push(opt_to_apply.infos[0].first_switch_info.otherwise_bb);
// new block that jumps to the correct discriminant case. This block is switched to if the discriminants are equal
let new_switch_data = BasicBlockData::new(Some(Terminator {
source_info: opt_to_apply.infos[0].second_switch_info.discr_source_info,
@@ -111,8 +113,7 @@ impl<'tcx> MirPass<'tcx> for EarlyOtherwiseBranch {
// the first and second discriminants are equal, so just pick one
discr: Operand::Copy(first_descriminant_place),
switch_ty: discr_type,
values: Cow::from(values_to_jump_to),
targets: targets_to_jump_to,
targets,
},
}));
@@ -176,7 +177,7 @@ struct SwitchDiscriminantInfo<'tcx> {
/// The basic block that the otherwise branch points to
otherwise_bb: BasicBlock,
/// Target along with the value being branched from. Otherwise is not included
targets_with_values: Vec<(BasicBlock, u128)>,
targets_with_values: Vec<(u128, BasicBlock)>,
discr_source_info: SourceInfo,
/// The place of the discriminant used in the switch
discr_used_in_switch: Place<'tcx>,
@@ -211,7 +212,7 @@ impl<'a, 'tcx> Helper<'a, 'tcx> {
let discr = self.find_switch_discriminant_info(bb, switch)?;
// go through each target, finding a discriminant read, and a switch
let results = discr.targets_with_values.iter().map(|(target, value)| {
let results = discr.targets_with_values.iter().map(|(value, target)| {
self.find_discriminant_switch_pairing(&discr, target.clone(), value.clone())
});
@@ -253,7 +254,7 @@ impl<'a, 'tcx> Helper<'a, 'tcx> {
}
// check that the value being matched on is the same. The
if this_bb_discr_info.targets_with_values.iter().find(|x| x.1 == value).is_none() {
if this_bb_discr_info.targets_with_values.iter().find(|x| x.0 == value).is_none() {
trace!("NO: values being matched on are not the same");
return None;
}
@@ -270,7 +271,7 @@ impl<'a, 'tcx> Helper<'a, 'tcx> {
// ```
// We check this by seeing that the value of the first discriminant is the only other discriminant value being used as a target in the second switch
if !(this_bb_discr_info.targets_with_values.len() == 1
&& this_bb_discr_info.targets_with_values[0].1 == value)
&& this_bb_discr_info.targets_with_values[0].0 == value)
{
trace!(
"NO: The second switch did not have only 1 target (besides otherwise) that had the same value as the value from the first switch that got us here"
@@ -296,18 +297,14 @@ impl<'a, 'tcx> Helper<'a, 'tcx> {
switch: &Terminator<'tcx>,
) -> Option<SwitchDiscriminantInfo<'tcx>> {
match &switch.kind {
TerminatorKind::SwitchInt { discr, targets, values, .. } => {
TerminatorKind::SwitchInt { discr, targets, .. } => {
let discr_local = discr.place()?.as_local()?;
// the declaration of the discriminant read. Place of this read is being used in the switch
let discr_decl = &self.body.local_decls()[discr_local];
let discr_ty = discr_decl.ty;
// the otherwise target lies as the last element
let otherwise_bb = targets.get(values.len())?.clone();
let targets_with_values = targets
.iter()
.zip(values.iter())
.map(|(t, v)| (t.clone(), v.clone()))
.collect();
let otherwise_bb = targets.otherwise();
let targets_with_values = targets.iter().collect();
// find the place of the adt where the discriminant is being read from
// assume this is the last statement of the block

View File

@@ -71,7 +71,6 @@ use rustc_middle::ty::GeneratorSubsts;
use rustc_middle::ty::{self, AdtDef, Ty, TyCtxt};
use rustc_target::abi::VariantIdx;
use rustc_target::spec::PanicStrategy;
use std::borrow::Cow;
use std::{iter, ops};
pub struct StateTransform;
@@ -839,11 +838,12 @@ fn insert_switch<'tcx>(
) {
let default_block = insert_term_block(body, default);
let (assign, discr) = transform.get_discr(body);
let switch_targets =
SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
let switch = TerminatorKind::SwitchInt {
discr: Operand::Move(discr),
switch_ty: transform.discr_ty,
values: Cow::from(cases.iter().map(|&(i, _)| i as u128).collect::<Vec<_>>()),
targets: cases.iter().map(|&(_, d)| d).chain(iter::once(default_block)).collect(),
targets: switch_targets,
};
let source_info = SourceInfo::outermost(body.span);

View File

@@ -771,7 +771,7 @@ impl<'a, 'tcx> MutVisitor<'tcx> for Integrator<'a, 'tcx> {
*target = self.update_target(*target);
}
TerminatorKind::SwitchInt { ref mut targets, .. } => {
for tgt in targets {
for tgt in targets.all_targets_mut() {
*tgt = self.update_target(*tgt);
}
}

View File

@@ -46,10 +46,13 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
discr: Operand::Copy(ref place) | Operand::Move(ref place),
switch_ty,
ref targets,
ref values,
..
} if targets.len() == 2 && values.len() == 1 && targets[0] != targets[1] => {
(place, values[0], switch_ty, targets[0], targets[1])
} if targets.iter().len() == 1 => {
let (value, target) = targets.iter().next().unwrap();
if target == targets.otherwise() {
continue;
}
(place, value, switch_ty, target, targets.otherwise())
}
// Only optimize switch int statements
_ => continue,

View File

@@ -29,17 +29,16 @@ impl<'tcx> MirPass<'tcx> for SimplifyBranches {
TerminatorKind::SwitchInt {
discr: Operand::Constant(ref c),
switch_ty,
ref values,
ref targets,
..
} => {
let constant = c.literal.try_eval_bits(tcx, param_env, switch_ty);
if let Some(constant) = constant {
let (otherwise, targets) = targets.split_last().unwrap();
let mut ret = TerminatorKind::Goto { target: *otherwise };
for (&v, t) in values.iter().zip(targets.iter()) {
let otherwise = targets.otherwise();
let mut ret = TerminatorKind::Goto { target: otherwise };
for (v, t) in targets.iter() {
if v == constant {
ret = TerminatorKind::Goto { target: *t };
ret = TerminatorKind::Goto { target: t };
break;
}
}

View File

@@ -1,8 +1,10 @@
use std::iter;
use super::MirPass;
use rustc_middle::{
mir::{
interpret::Scalar, BasicBlock, BinOp, Body, Operand, Place, Rvalue, Statement,
StatementKind, TerminatorKind,
StatementKind, SwitchTargets, TerminatorKind,
},
ty::{Ty, TyCtxt},
};
@@ -43,19 +45,21 @@ impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral {
Scalar::Ptr(_) => continue,
};
const FALSE: u128 = 0;
let mut new_targets = opt.targets.clone();
let first_is_false_target = opt.values[0] == FALSE;
let mut new_targets = opt.targets;
let first_value = new_targets.iter().next().unwrap().0;
let first_is_false_target = first_value == FALSE;
match opt.op {
BinOp::Eq => {
// if the assignment was Eq we want the true case to be first
if first_is_false_target {
new_targets.swap(0, 1);
new_targets.all_targets_mut().swap(0, 1);
}
}
BinOp::Ne => {
// if the assignment was Ne we want the false case to be first
if !first_is_false_target {
new_targets.swap(0, 1);
new_targets.all_targets_mut().swap(0, 1);
}
}
_ => unreachable!(),
@@ -96,7 +100,7 @@ impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral {
}
storage_deads_to_remove.push((stmt_idx, opt.bb_idx));
// if we have StorageDeads to remove then make sure to insert them at the top of each target
for bb_idx in new_targets.iter() {
for bb_idx in new_targets.all_targets() {
storage_deads_to_insert.push((
*bb_idx,
Statement {
@@ -107,13 +111,18 @@ impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral {
}
}
let terminator = bb.terminator_mut();
let [bb_cond, bb_otherwise] = match new_targets.all_targets() {
[a, b] => [*a, *b],
e => bug!("expected 2 switch targets, got: {:?}", e),
};
let targets = SwitchTargets::new(iter::once((new_value, bb_cond)), bb_otherwise);
let terminator = bb.terminator_mut();
terminator.kind = TerminatorKind::SwitchInt {
discr: Operand::Move(opt.to_switch_on),
switch_ty: opt.branch_value_ty,
values: vec![new_value].into(),
targets: new_targets,
targets,
};
}
@@ -138,15 +147,13 @@ impl<'a, 'tcx> OptimizationFinder<'a, 'tcx> {
.iter_enumerated()
.filter_map(|(bb_idx, bb)| {
// find switch
let (place_switched_on, values, targets, place_switched_on_moved) = match &bb
.terminator()
.kind
{
rustc_middle::mir::TerminatorKind::SwitchInt {
discr, values, targets, ..
} => Some((discr.place()?, values, targets, discr.is_move())),
_ => None,
}?;
let (place_switched_on, targets, place_switched_on_moved) =
match &bb.terminator().kind {
rustc_middle::mir::TerminatorKind::SwitchInt { discr, targets, .. } => {
Some((discr.place()?, targets, discr.is_move()))
}
_ => None,
}?;
// find the statement that assigns the place being switched on
bb.statements.iter().enumerate().rev().find_map(|(stmt_idx, stmt)| {
@@ -167,7 +174,6 @@ impl<'a, 'tcx> OptimizationFinder<'a, 'tcx> {
branch_value_scalar,
branch_value_ty,
op: *op,
values: values.clone().into_owned(),
targets: targets.clone(),
})
}
@@ -220,8 +226,6 @@ struct OptimizationInfo<'tcx> {
branch_value_ty: Ty<'tcx>,
/// Either Eq or Ne
op: BinOp,
/// Current values used in the switch target. This needs to be replaced with the branch_value
values: Vec<u128>,
/// Current targets used in the switch
targets: Vec<BasicBlock>,
targets: SwitchTargets<'tcx>,
}

View File

@@ -576,15 +576,13 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
.iter_enumerated()
.filter_map(|(bb_idx, bb)| {
let (discr_switched_on, targets_and_values) = match &bb.terminator().kind {
TerminatorKind::SwitchInt { targets, discr, values, .. } => {
// if values.len() == targets.len() - 1, we need to include None where no value is present
// such that the zip does not throw away targets. If no `otherwise` case is in targets, the zip will simply throw away the added None
let values_extended = values.iter().map(|x|Some(*x)).chain(once(None));
let targets_and_values:Vec<_> = targets.iter().zip(values_extended)
.map(|(target, value)| SwitchTargetAndValue{target:*target, value})
TerminatorKind::SwitchInt { targets, discr, .. } => {
let targets_and_values: Vec<_> = targets.iter()
.map(|(val, target)| SwitchTargetAndValue { target, value: Some(val) })
.chain(once(SwitchTargetAndValue { target: targets.otherwise(), value: None }))
.collect();
assert_eq!(targets.len(), targets_and_values.len());
(discr, targets_and_values)},
(discr, targets_and_values)
},
_ => return None,
};

View File

@@ -3,7 +3,8 @@
use crate::transform::MirPass;
use rustc_data_structures::stable_set::FxHashSet;
use rustc_middle::mir::{
BasicBlock, BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, TerminatorKind,
BasicBlock, BasicBlockData, Body, Local, Operand, Rvalue, StatementKind, SwitchTargets,
TerminatorKind,
};
use rustc_middle::ty::layout::TyAndLayout;
use rustc_middle::ty::{Ty, TyCtxt};
@@ -101,21 +102,15 @@ impl<'tcx> MirPass<'tcx> for UninhabitedEnumBranching {
trace!("allowed_variants = {:?}", allowed_variants);
if let TerminatorKind::SwitchInt { values, targets, .. } =
if let TerminatorKind::SwitchInt { targets, .. } =
&mut body.basic_blocks_mut()[bb].terminator_mut().kind
{
// take otherwise out early
let otherwise = targets.pop().unwrap();
assert_eq!(targets.len(), values.len());
let mut i = 0;
targets.retain(|_| {
let keep = allowed_variants.contains(&values[i]);
i += 1;
keep
});
targets.push(otherwise);
let new_targets = SwitchTargets::new(
targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
targets.otherwise(),
);
values.to_mut().retain(|var| allowed_variants.contains(var));
*targets = new_targets;
} else {
unreachable!()
}

View File

@@ -7,7 +7,6 @@ use crate::transform::MirPass;
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use rustc_middle::mir::*;
use rustc_middle::ty::TyCtxt;
use std::borrow::Cow;
pub struct UnreachablePropagation;
@@ -69,14 +68,15 @@ where
{
let terminator = match *terminator_kind {
TerminatorKind::Goto { target } if predicate(target) => TerminatorKind::Unreachable,
TerminatorKind::SwitchInt { ref discr, switch_ty, ref values, ref targets } => {
let original_targets_len = targets.len();
let (otherwise, targets) = targets.split_last().unwrap();
let (mut values, mut targets): (Vec<_>, Vec<_>) =
values.iter().zip(targets.iter()).filter(|(_, &t)| !predicate(t)).unzip();
TerminatorKind::SwitchInt { ref discr, switch_ty, ref targets } => {
let otherwise = targets.otherwise();
if !predicate(*otherwise) {
targets.push(*otherwise);
let original_targets_len = targets.iter().len() + 1;
let (mut values, mut targets): (Vec<_>, Vec<_>) =
targets.iter().filter(|(_, bb)| !predicate(*bb)).unzip();
if !predicate(otherwise) {
targets.push(otherwise);
} else {
values.pop();
}
@@ -91,8 +91,10 @@ where
TerminatorKind::SwitchInt {
discr: discr.clone(),
switch_ty,
values: Cow::from(values),
targets,
targets: SwitchTargets::new(
values.iter().copied().zip(targets.iter().copied()),
*targets.last().unwrap(),
),
}
} else {
return None;

View File

@@ -334,7 +334,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
TerminatorKind::Goto { target } => {
self.check_edge(location, *target, EdgeKind::Normal);
}
TerminatorKind::SwitchInt { targets, values, switch_ty, discr } => {
TerminatorKind::SwitchInt { targets, switch_ty, discr } => {
let ty = discr.ty(&self.body.local_decls, self.tcx);
if ty != *switch_ty {
self.fail(
@@ -345,19 +345,10 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
),
);
}
if targets.len() != values.len() + 1 {
self.fail(
location,
format!(
"encountered `SwitchInt` terminator with {} values, but {} targets (should be values+1)",
values.len(),
targets.len(),
),
);
}
for target in targets {
self.check_edge(location, *target, EdgeKind::Normal);
for (_, target) in targets.iter() {
self.check_edge(location, target, EdgeKind::Normal);
}
self.check_edge(location, targets.otherwise(), EdgeKind::Normal);
}
TerminatorKind::Drop { target, unwind, .. } => {
self.check_edge(location, *target, EdgeKind::Normal);