Support effects for particular edges of `SwitchInt`

2 files changed, 110 insertions, 7 deletions

diff --git a/src/librustc_mir/dataflow/generic/engine.rs b/src/librustc_mir/dataflow/generic/engine.rs
index b81f0adc201..d41ef2a392c 100644
--- a/src/librustc_mir/dataflow/generic/engine.rs
+++ b/src/librustc_mir/dataflow/generic/engine.rs
@@ -5,7 +5,7 @@ use std::fs;
 use std::path::PathBuf;
 
 use rustc::mir::{self, traversal, BasicBlock, Location};
-use rustc::ty::TyCtxt;
+use rustc::ty::{self, TyCtxt};
 use rustc_data_structures::work_queue::WorkQueue;
 use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::BitSet;
@@ -238,10 +238,15 @@ where
                 }
             }
 
-            SwitchInt { ref targets, .. } => {
-                for target in targets {
-                    self.propagate_bits_into_entry_set_for(in_out, *target, dirty_list);
-                }
+            SwitchInt { ref targets, ref values, ref discr, .. } => {
+                self.propagate_bits_into_switch_int_successors(
+                    in_out,
+                    (bb, bb_data),
+                    dirty_list,
+                    discr,
+                    &*values,
+                    &*targets,
+                );
             }
 
             Call { cleanup, ref destination, ref func, ref args, .. } => {
@@ -287,6 +292,66 @@ where
             dirty_queue.insert(bb);
         }
     }
+
+    fn propagate_bits_into_switch_int_successors(
+        &mut self,
+        in_out: &mut BitSet<A::Idx>,
+        (bb, bb_data): (BasicBlock, &mir::BasicBlockData<'tcx>),
+        dirty_list: &mut WorkQueue<BasicBlock>,
+        switch_on: &mir::Operand<'tcx>,
+        values: &[u128],
+        targets: &[BasicBlock],
+    ) {
+        match bb_data.statements.last().map(|stmt| &stmt.kind) {
+            // Look at the last statement to see if it is an assignment of an enum discriminant to
+            // the local that determines the target of a `SwitchInt` like so:
+            //   _42 = discriminant(..)
+            //   SwitchInt(_42, ..)
+            Some(mir::StatementKind::Assign(box (lhs, mir::Rvalue::Discriminant(enum_))))
+                if Some(lhs) == switch_on.place() =>
+            {
+                let adt = match enum_.ty(self.body, self.tcx).ty.kind {
+                    ty::Adt(def, _) => def,
+                    _ => bug!("Switch on discriminant of non-ADT"),
+                };
+
+                // MIR building adds discriminants to the `values` array in the same order as they
+                // are yielded by `AdtDef::discriminants`. We rely on this to match each
+                // discriminant in `values` to its corresponding variant in linear time.
+                let mut tmp = BitSet::new_empty(in_out.domain_size());
+                let mut discriminants = adt.discriminants(self.tcx);
+                for (value, target) in values.iter().zip(targets.iter().copied()) {
+                    let (variant_idx, _) =
+                        discriminants.find(|&(_, discr)| discr.val == *value).expect(
+                            "Order of `AdtDef::discriminants` differed \
+                                 from that of `SwitchInt::values`",
+                        );
+
+                    tmp.overwrite(in_out);
+                    self.analysis.apply_discriminant_switch_effect(
+                        &mut tmp,
+                        bb,
+                        enum_,
+                        adt,
+                        variant_idx,
+                    );
+                    self.propagate_bits_into_entry_set_for(&tmp, target, dirty_list);
+                }
+
+                std::mem::drop(tmp);
+
+                // Propagate dataflow state along the "otherwise" edge.
+                let otherwise = targets.last().copied().unwrap();
+                self.propagate_bits_into_entry_set_for(&in_out, otherwise, dirty_list);
+            }
+
+            _ => {
+                for target in targets.iter().copied() {
+                    self.propagate_bits_into_entry_set_for(&in_out, target, dirty_list);
+                }
+            }
+        }
+    }
 }
 
 // Graphviz
diff --git a/src/librustc_mir/dataflow/generic/mod.rs b/src/librustc_mir/dataflow/generic/mod.rs
index ea643042c5f..0f606240aeb 100644
--- a/src/librustc_mir/dataflow/generic/mod.rs
+++ b/src/librustc_mir/dataflow/generic/mod.rs
@@ -35,7 +35,8 @@
 use std::io;
 
 use rustc::mir::{self, BasicBlock, Location};
-use rustc::ty::TyCtxt;
+use rustc::ty::layout::VariantIdx;
+use rustc::ty::{self, TyCtxt};
 use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::{BitSet, HybridBitSet};
 use rustc_index::vec::{Idx, IndexVec};
@@ -172,7 +173,22 @@ pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
         return_place: &mir::Place<'tcx>,
     );
 
-    /// Calls the appropriate `Engine` constructor to find the fixpoint for this dataflow problem.
+    /// Updates the current dataflow state with the effect of taking a particular branch in a
+    /// `SwitchInt` terminator.
+    ///
+    /// Much like `apply_call_return_effect`, this effect is only propagated along a single
+    /// outgoing edge from this basic block.
+    fn apply_discriminant_switch_effect(
+        &self,
+        _state: &mut BitSet<Self::Idx>,
+        _block: BasicBlock,
+        _enum_place: &mir::Place<'tcx>,
+        _adt: &ty::AdtDef,
+        _variant: VariantIdx,
+    ) {
+    }
+
+    /// Creates an `Engine` to find the fixpoint for this dataflow problem.
     ///
     /// You shouldn't need to override this outside this module, since the combination of the
     /// default impl and the one for all `A: GenKillAnalysis` will do the right thing.
@@ -249,6 +265,17 @@ pub trait GenKillAnalysis<'tcx>: Analysis<'tcx> {
         args: &[mir::Operand<'tcx>],
         return_place: &mir::Place<'tcx>,
     );
+
+    /// See `Analysis::apply_discriminant_switch_effect`.
+    fn discriminant_switch_effect(
+        &self,
+        _state: &mut impl GenKill<Self::Idx>,
+        _block: BasicBlock,
+        _enum_place: &mir::Place<'tcx>,
+        _adt: &ty::AdtDef,
+        _variant: VariantIdx,
+    ) {
+    }
 }
 
 impl<A> Analysis<'tcx> for A
@@ -302,6 +329,17 @@ where
         self.call_return_effect(state, block, func, args, return_place);
     }
 
+    fn apply_discriminant_switch_effect(
+        &self,
+        state: &mut BitSet<Self::Idx>,
+        block: BasicBlock,
+        enum_place: &mir::Place<'tcx>,
+        adt: &ty::AdtDef,
+        variant: VariantIdx,
+    ) {
+        self.discriminant_switch_effect(state, block, enum_place, adt, variant);
+    }
+
     fn into_engine(
         self,
         tcx: TyCtxt<'tcx>,