Skip to content

Commit

Permalink
Fix #76803
Browse files Browse the repository at this point in the history 
Check that the variant index matches the target value from the SwitchInt we came from
  • Loading branch information
@simonvandel
simonvandel committed Sep 19, 2020
1 parent cad050b commit 738ed9b
Show file tree
Hide file tree
Showing 3 changed files with 55 additions and 44 deletions.
56 changes: 33 additions & 23 deletions compiler/rustc_mir/src/transform/simplify_try.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use rustc_middle::mir::visit::{NonUseContext, PlaceContext, Visitor};
use rustc_middle::mir::*;
use rustc_middle::ty::{self, List, Ty, TyCtxt};
use rustc_target::abi::VariantIdx;
use std::iter::{Enumerate, Peekable};
use std::iter::{once, Enumerate, Peekable};
use std::slice::Iter;

/// Simplifies arms of form `Variant(x) => Variant(x)` to just a move.
Expand Down Expand Up @@ -551,6 +551,12 @@ struct SimplifyBranchSameOptimization {
bb_to_opt_terminator: BasicBlock,
}

struct SwitchTargetAndValue {
target: BasicBlock,
// None in case of the `otherwise` case
value: Option<u128>,
}

struct SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
body: &'a Body<'tcx>,
tcx: TyCtxt<'tcx>,
Expand All @@ -562,8 +568,15 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
.basic_blocks()
.iter_enumerated()
.filter_map(|(bb_idx, bb)| {
let (discr_switched_on, targets) = match &bb.terminator().kind {
TerminatorKind::SwitchInt { targets, discr, .. } => (discr, targets),
let (discr_switched_on, targets_and_values):(_, Vec<_>) = 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 = targets.iter().zip(values_extended)
.map(|(target, value)| SwitchTargetAndValue{target:*target, value:value})
.collect();
(discr, targets_and_values)},
_ => return None,
};

Expand All @@ -587,9 +600,9 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
},
};

let mut iter_bbs_reachable = targets
let mut iter_bbs_reachable = targets_and_values
.iter()
.map(|idx| (*idx, &self.body.basic_blocks()[*idx]))
.map(|target_and_value| (target_and_value, &self.body.basic_blocks()[target_and_value.target]))
.filter(|(_, bb)| {
// Reaching `unreachable` is UB so assume it doesn't happen.
bb.terminator().kind != TerminatorKind::Unreachable
Expand All @@ -603,16 +616,16 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
})
.peekable();

let bb_first = iter_bbs_reachable.peek().map(|(idx, _)| *idx).unwrap_or(targets[0]);
let bb_first = iter_bbs_reachable.peek().map(|(idx, _)| *idx).unwrap_or(&targets_and_values[0]);
let mut all_successors_equivalent = StatementEquality::TrivialEqual;

// All successor basic blocks must be equal or contain statements that are pairwise considered equal.
for ((bb_l_idx,bb_l), (bb_r_idx,bb_r)) in iter_bbs_reachable.tuple_windows() {
for ((target_and_value_l,bb_l), (target_and_value_r,bb_r)) in iter_bbs_reachable.tuple_windows() {
let trivial_checks = bb_l.is_cleanup == bb_r.is_cleanup
&& bb_l.terminator().kind == bb_r.terminator().kind;
let statement_check = || {
bb_l.statements.iter().zip(&bb_r.statements).try_fold(StatementEquality::TrivialEqual, |acc,(l,r)| {
let stmt_equality = self.statement_equality(*adt_matched_on, &l, bb_l_idx, &r, bb_r_idx, self.tcx.sess.opts.debugging_opts.mir_opt_level);
let stmt_equality = self.statement_equality(*adt_matched_on, &l, target_and_value_l, &r, target_and_value_r);
if matches!(stmt_equality, StatementEquality::NotEqual) {
// short circuit
None
Expand All @@ -634,7 +647,7 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
// statements are trivially equal, so just take first
trace!("Statements are trivially equal");
Some(SimplifyBranchSameOptimization {
bb_to_goto: bb_first,
bb_to_goto: bb_first.target,
bb_to_opt_terminator: bb_idx,
})
}
Expand Down Expand Up @@ -669,10 +682,9 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
&self,
adt_matched_on: Place<'tcx>,
x: &Statement<'tcx>,
x_bb_idx: BasicBlock,
x_target_and_value: &SwitchTargetAndValue,
y: &Statement<'tcx>,
y_bb_idx: BasicBlock,
mir_opt_level: usize,
y_target_and_value: &SwitchTargetAndValue,
) -> StatementEquality {
let helper = |rhs: &Rvalue<'tcx>,
place: &Place<'tcx>,
Expand All @@ -691,13 +703,7 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {

match rhs {
Rvalue::Use(operand) if operand.place() == Some(adt_matched_on) => {
// FIXME(76803): This logic is currently broken because it does not take into
// account the current discriminant value.
if mir_opt_level > 2 {
StatementEquality::ConsideredEqual(side_to_choose)
} else {
StatementEquality::NotEqual
}
StatementEquality::ConsideredEqual(side_to_choose)
}
_ => {
trace!(
Expand All @@ -717,16 +723,20 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
(
StatementKind::Assign(box (_, rhs)),
StatementKind::SetDiscriminant { place, variant_index },
) => {
)
// we need to make sure that the switch value that targets the bb with SetDiscriminant (y), is the same as the variant index
if Some(variant_index.index() as u128) == y_target_and_value.value => {
// choose basic block of x, as that has the assign
helper(rhs, place, variant_index, x_bb_idx)
helper(rhs, place, variant_index, x_target_and_value.target)
}
(
StatementKind::SetDiscriminant { place, variant_index },
StatementKind::Assign(box (_, rhs)),
) => {
)
// we need to make sure that the switch value that targets the bb with SetDiscriminant (x), is the same as the variant index
if Some(variant_index.index() as u128) == x_target_and_value.value => {
// choose basic block of y, as that has the assign
helper(rhs, place, variant_index, y_bb_idx)
helper(rhs, place, variant_index, y_target_and_value.target)
}
_ => {
trace!("NO: statements `{:?}` and `{:?}` not considered equal", x, y);
Expand Down
16 changes: 7 additions & 9 deletions src/test/mir-opt/76803_regression.encode.SimplifyBranchSame.diff
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,22 +8,20 @@

bb0: {
_2 = discriminant(_1); // scope 0 at $DIR/76803_regression.rs:12:9: 12:16
- switchInt(move _2) -> [0_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/76803_regression.rs:12:9: 12:16
+ goto -> bb1; // scope 0 at $DIR/76803_regression.rs:12:9: 12:16
switchInt(move _2) -> [0_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/76803_regression.rs:12:9: 12:16
}

bb1: {
_0 = move _1; // scope 0 at $DIR/76803_regression.rs:13:14: 13:15
- goto -> bb3; // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
+ goto -> bb2; // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
goto -> bb3; // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
}

bb2: {
- discriminant(_0) = 1; // scope 0 at $DIR/76803_regression.rs:12:20: 12:27
- goto -> bb3; // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
- }
-
- bb3: {
discriminant(_0) = 1; // scope 0 at $DIR/76803_regression.rs:12:20: 12:27
goto -> bb3; // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
}
bb3: {
return; // scope 0 at $DIR/76803_regression.rs:15:2: 15:2
}
}
Expand Down
27 changes: 15 additions & 12 deletions src/test/mir-opt/simplify_arm.id.SimplifyBranchSame.diff
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,24 +13,27 @@

bb0: {
_2 = discriminant(_1); // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
switchInt(move _2) -> [0_isize: bb1, 1_isize: bb3, otherwise: bb2]; // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
- switchInt(move _2) -> [0_isize: bb1, 1_isize: bb3, otherwise: bb2]; // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
+ goto -> bb1; // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
}

bb1: {
discriminant(_0) = 0; // scope 0 at $DIR/simplify-arm.rs:12:17: 12:21
goto -> bb4; // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
}
bb2: {
unreachable; // scope 0 at $DIR/simplify-arm.rs:10:11: 10:12
}
bb3: {
- discriminant(_0) = 0; // scope 0 at $DIR/simplify-arm.rs:12:17: 12:21
- goto -> bb4; // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
- }
-
- bb2: {
- unreachable; // scope 0 at $DIR/simplify-arm.rs:10:11: 10:12
- }
-
- bb3: {
_0 = move _1; // scope 1 at $DIR/simplify-arm.rs:11:20: 11:27
goto -> bb4; // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
- goto -> bb4; // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
+ goto -> bb2; // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
}

bb4: {
- bb4: {
+ bb2: {
return; // scope 0 at $DIR/simplify-arm.rs:14:2: 14:2
}
}
Expand Down

0 comments on commit 738ed9b

Please sign in to comment.