Auto merge of #68528 - ecstatic-morse:maybe-init-variants, r=oli-obk

Mark other variants as uninitialized after switch on discriminant

During drop elaboration, which builds the drop ladder that handles destruction during stack unwinding, we attempt to remove MIR `Drop` terminators that will never be reached in practice. This reduces the number of basic blocks that are passed to LLVM, which should improve performance. In #66753, a user pointed out that unreachable `Drop` terminators are common in functions like `Option::unwrap`, which move out of an `enum`. While discussing possible remedies for that issue, @eddyb suggested moving const-checking after drop elaboration. This would allow the former, which looks for `Drop` terminators and replicates a small amount of drop elaboration to determine whether a dropped local has been moved out, leverage the work done by the latter.

However, it turns out that drop elaboration is not as precise as it could be when it comes to eliminating useless drop terminators. For example, let's look at the code for `unwrap_or`.

```rust
fn unwrap_or<T>(opt: Option<T>, default: T) -> T {
    match opt {
        Some(inner) => inner,
        None => default,
    }
}
```

`opt` never needs to be dropped, since it is either moved out of (if it is `Some`) or has no drop glue (if it is `None`), and `default` only needs to be dropped if `opt` is `Some`. This is not reflected in the MIR we currently pass to codegen.

![pasted_image](https://user-images.githubusercontent.com/29463364/73384403-109a0d80-4280-11ea-8500-0637b368f2dc.png)

@eddyb also suggested the solution to this problem. When we switch on an enum discriminant, we should be marking all fields in other variants as definitely uninitialized. I implemented this on top of alongside a small optimization (split out into #68943) that suppresses drop terminators for enum variants with no fields (e.g. `Option::None`). This is the resulting MIR for `unwrap_or`.

![after](https://user-images.githubusercontent.com/29463364/73384823-e432c100-4280-11ea-84bd-d0bcc3b777b4.png)

In concert with #68943, this change speeds up many [optimized and debug builds](https://perf.rust-lang.org/compare.html?start=d55f3e9f1da631c636b54a7c22c1caccbe4bf0db&end=0077a7aa11ebc2462851676f9f464d5221b17d6a). We need to carefully investigate whether I have introduced any miscompilations before merging this. Code that never drops anything would be very fast indeed until memory is exhausted.

src/librustc/mir/mod.rs

@@ -2046,6 +2046,15 @@ impl<'tcx> Operand<'tcx> {
             Operand::Move(place) => Operand::Copy(place),
         }
     }
+
+    /// Returns the `Place` that is the target of this `Operand`, or `None` if this `Operand` is a
+    /// constant.
+    pub fn place(&self) -> Option<&Place<'tcx>> {
+        match self {
+            Operand::Copy(place) | Operand::Move(place) => Some(place),
+            Operand::Constant(_) => None,
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////

src/librustc_mir/borrow_check/diagnostics/conflict_errors.rs

@@ -1340,7 +1340,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                 // there.
                 let mut mpis = vec![mpi];
                 let move_paths = &self.move_data.move_paths;
-                mpis.extend(move_paths[mpi].parents(move_paths));
+                mpis.extend(move_paths[mpi].parents(move_paths).map(|(mpi, _)| mpi));
 
                 for moi in &self.move_data.loc_map[location] {
                     debug!("report_use_of_moved_or_uninitialized: moi={:?}", moi);

src/librustc_mir/borrow_check/mod.rs

@@ -1599,9 +1599,9 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
     ) {
         if let Some(mpi) = self.move_path_for_place(place_span.0) {
             let move_paths = &self.move_data.move_paths;
-            let mut child = move_paths[mpi].first_child;
-            while let Some(child_mpi) = child {
-                let child_move_path = &move_paths[child_mpi];
+
+            let root_path = &move_paths[mpi];
+            for (child_mpi, child_move_path) in root_path.children(move_paths) {
                 let last_proj = child_move_path.place.projection.last().unwrap();
                 if let ProjectionElem::ConstantIndex { offset, from_end, .. } = last_proj {
                     debug_assert!(!from_end, "Array constant indexing shouldn't be `from_end`.");
@@ -1623,7 +1623,6 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                         }
                     }
                 }
-                child = child_move_path.next_sibling;
             }
         }
     }

src/librustc_mir/borrow_check/nll.rs

@@ -92,7 +92,7 @@ fn populate_polonius_move_facts(
     for (child, move_path) in move_data.move_paths.iter_enumerated() {
         all_facts
             .child
-            .extend(move_path.parents(&move_data.move_paths).iter().map(|&parent| (child, parent)));
+            .extend(move_path.parents(&move_data.move_paths).map(|(parent, _)| (child, parent)));
     }
 
     // initialized_at

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

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>,

src/librustc_mir/dataflow/impls/mod.rs

@@ -3,7 +3,8 @@
 //! zero-sized structure.
 
 use rustc::mir::{self, Body, Location};
-use rustc::ty::TyCtxt;
+use rustc::ty::layout::VariantIdx;
+use rustc::ty::{self, TyCtxt};
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::Idx;
 
@@ -12,12 +13,13 @@ use super::MoveDataParamEnv;
 use crate::util::elaborate_drops::DropFlagState;
 
 use super::generic::{AnalysisDomain, GenKill, GenKillAnalysis};
-use super::move_paths::{HasMoveData, InitIndex, InitKind, MoveData, MovePathIndex};
+use super::move_paths::{HasMoveData, InitIndex, InitKind, LookupResult, MoveData, MovePathIndex};
 use super::{BottomValue, GenKillSet};
 
 use super::drop_flag_effects_for_function_entry;
 use super::drop_flag_effects_for_location;
 use super::on_lookup_result_bits;
+use crate::dataflow::drop_flag_effects;
 
 mod borrowed_locals;
 mod storage_liveness;
@@ -336,6 +338,37 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
             },
         );
     }
+
+    fn discriminant_switch_effect(
+        &self,
+        trans: &mut impl GenKill<Self::Idx>,
+        _block: mir::BasicBlock,
+        enum_place: &mir::Place<'tcx>,
+        _adt: &ty::AdtDef,
+        variant: VariantIdx,
+    ) {
+        let enum_mpi = match self.move_data().rev_lookup.find(enum_place.as_ref()) {
+            LookupResult::Exact(mpi) => mpi,
+            LookupResult::Parent(_) => return,
+        };
+
+        // Kill all move paths that correspond to variants other than this one
+        let move_paths = &self.move_data().move_paths;
+        let enum_path = &move_paths[enum_mpi];
+        for (mpi, variant_path) in enum_path.children(move_paths) {
+            trans.kill(mpi);
+            match variant_path.place.projection.last().unwrap() {
+                mir::ProjectionElem::Downcast(_, idx) if *idx == variant => continue,
+                _ => drop_flag_effects::on_all_children_bits(
+                    self.tcx,
+                    self.body,
+                    self.move_data(),
+                    mpi,
+                    |mpi| trans.kill(mpi),
+                ),
+            }
+        }
+    }
 }
 
 impl<'tcx> AnalysisDomain<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {

src/librustc_mir/dataflow/move_paths/mod.rs

@@ -58,19 +58,32 @@ pub struct MovePath<'tcx> {
 }
 
 impl<'tcx> MovePath<'tcx> {
-    pub fn parents(
+    /// Returns an iterator over the parents of `self`.
+    pub fn parents<'a>(
         &self,
-        move_paths: &IndexVec<MovePathIndex, MovePath<'_>>,
-    ) -> Vec<MovePathIndex> {
-        let mut parents = Vec::new();
-
-        let mut curr_parent = self.parent;
-        while let Some(parent_mpi) = curr_parent {
-            parents.push(parent_mpi);
-            curr_parent = move_paths[parent_mpi].parent;
+        move_paths: &'a IndexVec<MovePathIndex, MovePath<'tcx>>,
+    ) -> impl 'a + Iterator<Item = (MovePathIndex, &'a MovePath<'tcx>)> {
+        let first = self.parent.map(|mpi| (mpi, &move_paths[mpi]));
+        MovePathLinearIter {
+            next: first,
+            fetch_next: move |_, parent: &MovePath<'_>| {
+                parent.parent.map(|mpi| (mpi, &move_paths[mpi]))
+            },
         }
+    }
 
-        parents
+    /// Returns an iterator over the immediate children of `self`.
+    pub fn children<'a>(
+        &self,
+        move_paths: &'a IndexVec<MovePathIndex, MovePath<'tcx>>,
+    ) -> impl 'a + Iterator<Item = (MovePathIndex, &'a MovePath<'tcx>)> {
+        let first = self.first_child.map(|mpi| (mpi, &move_paths[mpi]));
+        MovePathLinearIter {
+            next: first,
+            fetch_next: move |_, child: &MovePath<'_>| {
+                child.next_sibling.map(|mpi| (mpi, &move_paths[mpi]))
+            },
+        }
     }
 
     /// Finds the closest descendant of `self` for which `f` returns `true` using a breadth-first
@@ -131,6 +144,25 @@ impl<'tcx> fmt::Display for MovePath<'tcx> {
     }
 }
 
+#[allow(unused)]
+struct MovePathLinearIter<'a, 'tcx, F> {
+    next: Option<(MovePathIndex, &'a MovePath<'tcx>)>,
+    fetch_next: F,
+}
+
+impl<'a, 'tcx, F> Iterator for MovePathLinearIter<'a, 'tcx, F>
+where
+    F: FnMut(MovePathIndex, &'a MovePath<'tcx>) -> Option<(MovePathIndex, &'a MovePath<'tcx>)>,
+{
+    type Item = (MovePathIndex, &'a MovePath<'tcx>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let ret = self.next.take()?;
+        self.next = (self.fetch_next)(ret.0, ret.1);
+        Some(ret)
+    }
+}
+
 #[derive(Debug)]
 pub struct MoveData<'tcx> {
     pub move_paths: IndexVec<MovePathIndex, MovePath<'tcx>>,