Refactor how SwitchInt stores jump targets

10 files changed, 86 insertions, 97 deletions

diff --git a/compiler/rustc_mir/src/transform/early_otherwise_branch.rs b/compiler/rustc_mir/src/transform/early_otherwise_branch.rs
index 237a5c7864b..f97dcf4852d 100644
--- a/compiler/rustc_mir/src/transform/early_otherwise_branch.rs
+++ b/compiler/rustc_mir/src/transform/early_otherwise_branch.rs
@@ -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
diff --git a/compiler/rustc_mir/src/transform/generator.rs b/compiler/rustc_mir/src/transform/generator.rs
index 924bb4996fc..039d4753a8c 100644
--- a/compiler/rustc_mir/src/transform/generator.rs
+++ b/compiler/rustc_mir/src/transform/generator.rs
@@ -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);
diff --git a/compiler/rustc_mir/src/transform/inline.rs b/compiler/rustc_mir/src/transform/inline.rs
index bec1eb79047..34aaefdcbea 100644
--- a/compiler/rustc_mir/src/transform/inline.rs
+++ b/compiler/rustc_mir/src/transform/inline.rs
@@ -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);
                 }
             }
diff --git a/compiler/rustc_mir/src/transform/match_branches.rs b/compiler/rustc_mir/src/transform/match_branches.rs
index dad3812c5cd..8b2d6b09aa8 100644
--- a/compiler/rustc_mir/src/transform/match_branches.rs
+++ b/compiler/rustc_mir/src/transform/match_branches.rs
@@ -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,
diff --git a/compiler/rustc_mir/src/transform/simplify_branches.rs b/compiler/rustc_mir/src/transform/simplify_branches.rs
index 161856a38ee..5f63c03993d 100644
--- a/compiler/rustc_mir/src/transform/simplify_branches.rs
+++ b/compiler/rustc_mir/src/transform/simplify_branches.rs
@@ -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;
                             }
                         }
diff --git a/compiler/rustc_mir/src/transform/simplify_comparison_integral.rs b/compiler/rustc_mir/src/transform/simplify_comparison_integral.rs
index 9f837cf78a6..50f5fbb3bc2 100644
--- a/compiler/rustc_mir/src/transform/simplify_comparison_integral.rs
+++ b/compiler/rustc_mir/src/transform/simplify_comparison_integral.rs
@@ -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>,
 }
diff --git a/compiler/rustc_mir/src/transform/simplify_try.rs b/compiler/rustc_mir/src/transform/simplify_try.rs
index a4e7a5a9453..27bb1def726 100644
--- a/compiler/rustc_mir/src/transform/simplify_try.rs
+++ b/compiler/rustc_mir/src/transform/simplify_try.rs
@@ -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,
                 };
 
diff --git a/compiler/rustc_mir/src/transform/uninhabited_enum_branching.rs b/compiler/rustc_mir/src/transform/uninhabited_enum_branching.rs
index 87906e83ed5..465832c89fd 100644
--- a/compiler/rustc_mir/src/transform/uninhabited_enum_branching.rs
+++ b/compiler/rustc_mir/src/transform/uninhabited_enum_branching.rs
@@ -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);
-
-                values.to_mut().retain(|var| allowed_variants.contains(var));
+                let new_targets = SwitchTargets::new(
+                    targets.iter().filter(|(val, _)| allowed_variants.contains(val)),
+                    targets.otherwise(),
+                );
+
+                *targets = new_targets;
             } else {
                 unreachable!()
             }
diff --git a/compiler/rustc_mir/src/transform/unreachable_prop.rs b/compiler/rustc_mir/src/transform/unreachable_prop.rs
index c6426a06ea1..f6d39dae342 100644
--- a/compiler/rustc_mir/src/transform/unreachable_prop.rs
+++ b/compiler/rustc_mir/src/transform/unreachable_prop.rs
@@ -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();
+        TerminatorKind::SwitchInt { ref discr, switch_ty, ref targets } => {
+            let otherwise = targets.otherwise();
+
+            let original_targets_len = targets.iter().len() + 1;
             let (mut values, mut targets): (Vec<_>, Vec<_>) =
-                values.iter().zip(targets.iter()).filter(|(_, &t)| !predicate(t)).unzip();
+                targets.iter().filter(|(_, bb)| !predicate(*bb)).unzip();
 
-            if !predicate(*otherwise) {
-                targets.push(*otherwise);
+            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;
diff --git a/compiler/rustc_mir/src/transform/validate.rs b/compiler/rustc_mir/src/transform/validate.rs
index cf51e86c5bc..48477f60ef7 100644
--- a/compiler/rustc_mir/src/transform/validate.rs
+++ b/compiler/rustc_mir/src/transform/validate.rs
@@ -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);