Auto merge of #77796 - jonas-schievink:switchint-refactor, r=oli-obk

Refactor how SwitchInt stores jump targets

Closes #65693

Author: bors

compiler/rustc_codegen_ssa/src/mir/block.rs

@@ -12,9 +12,9 @@ use crate::MemFlags;
 use rustc_ast as ast;
 use rustc_hir::lang_items::LangItem;
 use rustc_index::vec::Idx;
-use rustc_middle::mir;
 use rustc_middle::mir::interpret::ConstValue;
 use rustc_middle::mir::AssertKind;
+use rustc_middle::mir::{self, SwitchTargets};
 use rustc_middle::ty::layout::{FnAbiExt, HasTyCtxt};
 use rustc_middle::ty::print::with_no_trimmed_paths;
 use rustc_middle::ty::{self, Instance, Ty, TypeFoldable};
@@ -24,8 +24,6 @@ use rustc_target::abi::call::{ArgAbi, FnAbi, PassMode};
 use rustc_target::abi::{self, LayoutOf};
 use rustc_target::spec::abi::Abi;
 
-use std::borrow::Cow;
-
 /// Used by `FunctionCx::codegen_terminator` for emitting common patterns
 /// e.g., creating a basic block, calling a function, etc.
 struct TerminatorCodegenHelper<'tcx> {
@@ -198,42 +196,37 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
         mut bx: Bx,
         discr: &mir::Operand<'tcx>,
         switch_ty: Ty<'tcx>,
-        values: &Cow<'tcx, [u128]>,
-        targets: &Vec<mir::BasicBlock>,
+        targets: &SwitchTargets,
     ) {
         let discr = self.codegen_operand(&mut bx, &discr);
         // `switch_ty` is redundant, sanity-check that.
         assert_eq!(discr.layout.ty, switch_ty);
-        if targets.len() == 2 {
-            // If there are two targets, emit br instead of switch
-            let lltrue = helper.llblock(self, targets[0]);
-            let llfalse = helper.llblock(self, targets[1]);
+        helper.maybe_sideeffect(self.mir, &mut bx, targets.all_targets());
+
+        let mut target_iter = targets.iter();
+        if target_iter.len() == 1 {
+            // If there are two targets (one conditional, one fallback), emit br instead of switch
+            let (test_value, target) = target_iter.next().unwrap();
+            let lltrue = helper.llblock(self, target);
+            let llfalse = helper.llblock(self, targets.otherwise());
             if switch_ty == bx.tcx().types.bool {
-                helper.maybe_sideeffect(self.mir, &mut bx, targets.as_slice());
                 // Don't generate trivial icmps when switching on bool
-                if let [0] = values[..] {
-                    bx.cond_br(discr.immediate(), llfalse, lltrue);
-                } else {
-                    assert_eq!(&values[..], &[1]);
-                    bx.cond_br(discr.immediate(), lltrue, llfalse);
+                match test_value {
+                    0 => bx.cond_br(discr.immediate(), llfalse, lltrue),
+                    1 => bx.cond_br(discr.immediate(), lltrue, llfalse),
+                    _ => bug!(),
                 }
             } else {
                 let switch_llty = bx.immediate_backend_type(bx.layout_of(switch_ty));
-                let llval = bx.const_uint_big(switch_llty, values[0]);
+                let llval = bx.const_uint_big(switch_llty, test_value);
                 let cmp = bx.icmp(IntPredicate::IntEQ, discr.immediate(), llval);
-                helper.maybe_sideeffect(self.mir, &mut bx, targets.as_slice());
                 bx.cond_br(cmp, lltrue, llfalse);
             }
         } else {
-            helper.maybe_sideeffect(self.mir, &mut bx, targets.as_slice());
-            let (otherwise, targets) = targets.split_last().unwrap();
             bx.switch(
                 discr.immediate(),
-                helper.llblock(self, *otherwise),
-                values
-                    .iter()
-                    .zip(targets)
-                    .map(|(&value, target)| (value, helper.llblock(self, *target))),
+                helper.llblock(self, targets.otherwise()),
+                target_iter.map(|(value, target)| (value, helper.llblock(self, target))),
             );
         }
     }
@@ -975,8 +968,8 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 helper.funclet_br(self, &mut bx, target);
             }
 
-            mir::TerminatorKind::SwitchInt { ref discr, switch_ty, ref values, ref targets } => {
-                self.codegen_switchint_terminator(helper, bx, discr, switch_ty, values, targets);
+            mir::TerminatorKind::SwitchInt { ref discr, switch_ty, ref targets } => {
+                self.codegen_switchint_terminator(helper, bx, discr, switch_ty, targets);
             }
 
             mir::TerminatorKind::Return => {

compiler/rustc_middle/src/mir/terminator/mod.rs

@@ -1,6 +1,7 @@
 use crate::mir::interpret::Scalar;
 use crate::ty::{self, Ty, TyCtxt};
 use rustc_ast::{InlineAsmOptions, InlineAsmTemplatePiece};
+use smallvec::{smallvec, SmallVec};
 
 use super::{
     AssertMessage, BasicBlock, InlineAsmOperand, Operand, Place, SourceInfo, Successors,
@@ -16,6 +17,87 @@ use std::slice;
 
 pub use super::query::*;
 
+#[derive(Debug, Clone, TyEncodable, TyDecodable, HashStable, PartialEq)]
+pub struct SwitchTargets {
+    /// Possible values. The locations to branch to in each case
+    /// are found in the corresponding indices from the `targets` vector.
+    values: SmallVec<[u128; 1]>,
+
+    /// Possible branch sites. The last element of this vector is used
+    /// for the otherwise branch, so targets.len() == values.len() + 1
+    /// should hold.
+    //
+    // This invariant is quite non-obvious and also could be improved.
+    // One way to make this invariant is to have something like this instead:
+    //
+    // branches: Vec<(ConstInt, BasicBlock)>,
+    // otherwise: Option<BasicBlock> // exhaustive if None
+    //
+    // However we’ve decided to keep this as-is until we figure a case
+    // where some other approach seems to be strictly better than other.
+    targets: SmallVec<[BasicBlock; 2]>,
+}
+
+impl SwitchTargets {
+    /// Creates switch targets from an iterator of values and target blocks.
+    ///
+    /// The iterator may be empty, in which case the `SwitchInt` instruction is equivalent to
+    /// `goto otherwise;`.
+    pub fn new(targets: impl Iterator<Item = (u128, BasicBlock)>, otherwise: BasicBlock) -> Self {
+        let (values, mut targets): (SmallVec<_>, SmallVec<_>) = targets.unzip();
+        targets.push(otherwise);
+        Self { values: values.into(), targets }
+    }
+
+    /// Builds a switch targets definition that jumps to `then` if the tested value equals `value`,
+    /// and to `else_` if not.
+    pub fn static_if(value: u128, then: BasicBlock, else_: BasicBlock) -> Self {
+        Self { values: smallvec![value], targets: smallvec![then, else_] }
+    }
+
+    /// Returns the fallback target that is jumped to when none of the values match the operand.
+    pub fn otherwise(&self) -> BasicBlock {
+        *self.targets.last().unwrap()
+    }
+
+    /// Returns an iterator over the switch targets.
+    ///
+    /// The iterator will yield tuples containing the value and corresponding target to jump to, not
+    /// including the `otherwise` fallback target.
+    ///
+    /// Note that this may yield 0 elements. Only the `otherwise` branch is mandatory.
+    pub fn iter(&self) -> SwitchTargetsIter<'_> {
+        SwitchTargetsIter { inner: self.values.iter().zip(self.targets.iter()) }
+    }
+
+    /// Returns a slice with all possible jump targets (including the fallback target).
+    pub fn all_targets(&self) -> &[BasicBlock] {
+        &self.targets
+    }
+
+    pub fn all_targets_mut(&mut self) -> &mut [BasicBlock] {
+        &mut self.targets
+    }
+}
+
+pub struct SwitchTargetsIter<'a> {
+    inner: iter::Zip<slice::Iter<'a, u128>, slice::Iter<'a, BasicBlock>>,
+}
+
+impl<'a> Iterator for SwitchTargetsIter<'a> {
+    type Item = (u128, BasicBlock);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.inner.next().map(|(val, bb)| (*val, *bb))
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.inner.size_hint()
+    }
+}
+
+impl<'a> ExactSizeIterator for SwitchTargetsIter<'a> {}
+
 #[derive(Clone, TyEncodable, TyDecodable, HashStable, PartialEq)]
 pub enum TerminatorKind<'tcx> {
     /// Block should have one successor in the graph; we jump there.
@@ -32,23 +114,7 @@ pub enum TerminatorKind<'tcx> {
         /// FIXME: remove this redundant information. Currently, it is relied on by pretty-printing.
         switch_ty: Ty<'tcx>,
 
-        /// Possible values. The locations to branch to in each case
-        /// are found in the corresponding indices from the `targets` vector.
-        values: Cow<'tcx, [u128]>,
-
-        /// Possible branch sites. The last element of this vector is used
-        /// for the otherwise branch, so targets.len() == values.len() + 1
-        /// should hold.
-        //
-        // This invariant is quite non-obvious and also could be improved.
-        // One way to make this invariant is to have something like this instead:
-        //
-        // branches: Vec<(ConstInt, BasicBlock)>,
-        // otherwise: Option<BasicBlock> // exhaustive if None
-        //
-        // However we’ve decided to keep this as-is until we figure a case
-        // where some other approach seems to be strictly better than other.
-        targets: Vec<BasicBlock>,
+        targets: SwitchTargets,
     },
 
     /// Indicates that the landing pad is finished and unwinding should
@@ -227,12 +293,10 @@ impl<'tcx> TerminatorKind<'tcx> {
         t: BasicBlock,
         f: BasicBlock,
     ) -> TerminatorKind<'tcx> {
-        static BOOL_SWITCH_FALSE: &[u128] = &[0];
         TerminatorKind::SwitchInt {
             discr: cond,
             switch_ty: tcx.types.bool,
-            values: From::from(BOOL_SWITCH_FALSE),
-            targets: vec![f, t],
+            targets: SwitchTargets::static_if(0, f, t),
         }
     }
 
@@ -263,7 +327,7 @@ impl<'tcx> TerminatorKind<'tcx> {
             | FalseUnwind { real_target: ref t, unwind: Some(ref u) } => {
                 Some(t).into_iter().chain(slice::from_ref(u))
             }
-            SwitchInt { ref targets, .. } => None.into_iter().chain(&targets[..]),
+            SwitchInt { ref targets, .. } => None.into_iter().chain(&targets.targets[..]),
             FalseEdge { ref real_target, ref imaginary_target } => {
                 Some(real_target).into_iter().chain(slice::from_ref(imaginary_target))
             }
@@ -297,7 +361,7 @@ impl<'tcx> TerminatorKind<'tcx> {
             | FalseUnwind { real_target: ref mut t, unwind: Some(ref mut u) } => {
                 Some(t).into_iter().chain(slice::from_mut(u))
             }
-            SwitchInt { ref mut targets, .. } => None.into_iter().chain(&mut targets[..]),
+            SwitchInt { ref mut targets, .. } => None.into_iter().chain(&mut targets.targets[..]),
             FalseEdge { ref mut real_target, ref mut imaginary_target } => {
                 Some(real_target).into_iter().chain(slice::from_mut(imaginary_target))
             }
@@ -469,11 +533,12 @@ impl<'tcx> TerminatorKind<'tcx> {
         match *self {
             Return | Resume | Abort | Unreachable | GeneratorDrop => vec![],
             Goto { .. } => vec!["".into()],
-            SwitchInt { ref values, switch_ty, .. } => ty::tls::with(|tcx| {
+            SwitchInt { ref targets, switch_ty, .. } => ty::tls::with(|tcx| {
                 let param_env = ty::ParamEnv::empty();
                 let switch_ty = tcx.lift(&switch_ty).unwrap();
                 let size = tcx.layout_of(param_env.and(switch_ty)).unwrap().size;
-                values
+                targets
+                    .values
                     .iter()
                     .map(|&u| {
                         ty::Const::from_scalar(tcx, Scalar::from_uint(u, size), switch_ty)

compiler/rustc_middle/src/mir/type_foldable.rs

@@ -21,10 +21,9 @@ impl<'tcx> TypeFoldable<'tcx> for Terminator<'tcx> {
 
         let kind = match self.kind {
             Goto { target } => Goto { target },
-            SwitchInt { ref discr, switch_ty, ref values, ref targets } => SwitchInt {
+            SwitchInt { ref discr, switch_ty, ref targets } => SwitchInt {
                 discr: discr.fold_with(folder),
                 switch_ty: switch_ty.fold_with(folder),
-                values: values.clone(),
                 targets: targets.clone(),
             },
             Drop { ref place, target, unwind } => {

compiler/rustc_middle/src/mir/visit.rs

@@ -453,7 +453,6 @@ macro_rules! make_mir_visitor {
                     TerminatorKind::SwitchInt {
                         discr,
                         switch_ty,
-                        values: _,
                         targets: _
                     } => {
                         self.visit_operand(discr, location);

compiler/rustc_mir/src/borrow_check/invalidation.rs

@@ -117,7 +117,7 @@ impl<'cx, 'tcx> Visitor<'tcx> for InvalidationGenerator<'cx, 'tcx> {
         self.check_activations(location);
 
         match &terminator.kind {
-            TerminatorKind::SwitchInt { ref discr, switch_ty: _, values: _, targets: _ } => {
+            TerminatorKind::SwitchInt { ref discr, switch_ty: _, targets: _ } => {
                 self.consume_operand(location, discr);
             }
             TerminatorKind::Drop { place: drop_place, target: _, unwind: _ } => {

compiler/rustc_mir/src/borrow_check/mod.rs

@@ -671,7 +671,7 @@ impl<'cx, 'tcx> dataflow::ResultsVisitor<'cx, 'tcx> for MirBorrowckCtxt<'cx, 'tc
         self.check_activations(loc, span, flow_state);
 
         match term.kind {
-            TerminatorKind::SwitchInt { ref discr, switch_ty: _, values: _, targets: _ } => {
+            TerminatorKind::SwitchInt { ref discr, switch_ty: _, targets: _ } => {
                 self.consume_operand(loc, (discr, span), flow_state);
             }
             TerminatorKind::Drop { place: ref drop_place, target: _, unwind: _ } => {

compiler/rustc_mir/src/borrow_check/type_check/mod.rs

@@ -1777,7 +1777,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 self.assert_iscleanup(body, block_data, target, is_cleanup)
             }
             TerminatorKind::SwitchInt { ref targets, .. } => {
-                for target in targets {
+                for target in targets.all_targets() {
                     self.assert_iscleanup(body, block_data, *target, is_cleanup);
                 }
             }

compiler/rustc_mir/src/dataflow/framework/direction.rs

@@ -1,5 +1,5 @@
 use rustc_index::bit_set::BitSet;
-use rustc_middle::mir::{self, BasicBlock, Location};
+use rustc_middle::mir::{self, BasicBlock, Location, SwitchTargets};
 use rustc_middle::ty::TyCtxt;
 use std::ops::RangeInclusive;
 
@@ -488,11 +488,10 @@ impl Direction for Forward {
                 }
             }
 
-            SwitchInt { ref targets, ref values, ref discr, switch_ty: _ } => {
+            SwitchInt { ref targets, ref discr, switch_ty: _ } => {
                 let mut applier = SwitchIntEdgeEffectApplier {
                     exit_state,
-                    targets: targets.as_ref(),
-                    values: values.as_ref(),
+                    targets,
                     propagate,
                     effects_applied: false,
                 };
@@ -504,8 +503,8 @@ impl Direction for Forward {
                 } = applier;
 
                 if !effects_applied {
-                    for &target in targets.iter() {
-                        propagate(target, exit_state);
+                    for target in targets.all_targets() {
+                        propagate(*target, exit_state);
                     }
                 }
             }
@@ -515,8 +514,7 @@ impl Direction for Forward {
 
 struct SwitchIntEdgeEffectApplier<'a, D, F> {
     exit_state: &'a mut D,
-    values: &'a [u128],
-    targets: &'a [BasicBlock],
+    targets: &'a SwitchTargets,
     propagate: F,
 
     effects_applied: bool,
@@ -531,15 +529,15 @@ where
         assert!(!self.effects_applied);
 
         let mut tmp = None;
-        for (&value, &target) in self.values.iter().zip(self.targets.iter()) {
+        for (value, target) in self.targets.iter() {
             let tmp = opt_clone_from_or_clone(&mut tmp, self.exit_state);
             apply_edge_effect(tmp, SwitchIntTarget { value: Some(value), target });
             (self.propagate)(target, tmp);
         }
 
         // Once we get to the final, "otherwise" branch, there is no need to preserve `exit_state`,
         // so pass it directly to `apply_edge_effect` to save a clone of the dataflow state.
-        let otherwise = self.targets.last().copied().unwrap();
+        let otherwise = self.targets.otherwise();
         apply_edge_effect(self.exit_state, SwitchIntTarget { value: None, target: otherwise });
         (self.propagate)(otherwise, self.exit_state);
 

compiler/rustc_mir/src/interpret/terminator.rs

@@ -24,16 +24,16 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
 
             Goto { target } => self.go_to_block(target),
 
-            SwitchInt { ref discr, ref values, ref targets, switch_ty } => {
+            SwitchInt { ref discr, ref targets, switch_ty } => {
                 let discr = self.read_immediate(self.eval_operand(discr, None)?)?;
                 trace!("SwitchInt({:?})", *discr);
                 assert_eq!(discr.layout.ty, switch_ty);
 
                 // Branch to the `otherwise` case by default, if no match is found.
-                assert!(!targets.is_empty());
-                let mut target_block = targets[targets.len() - 1];
+                assert!(!targets.iter().is_empty());
+                let mut target_block = targets.otherwise();
 
-                for (index, &const_int) in values.iter().enumerate() {
+                for (const_int, target) in targets.iter() {
                     // Compare using binary_op, to also support pointer values
                     let res = self
                         .overflowing_binary_op(
@@ -43,7 +43,7 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
                         )?
                         .0;
                     if res.to_bool()? {
-                        target_block = targets[index];
+                        target_block = target;
                         break;
                     }
                 }