Auto merge of #77242 - ecstatic-morse:dataflow-switch-int, r=jonas-schievink

Replace `discriminant_switch_effect` with more general version

#68528 added a new edge-specific effect for `SwitchInt` terminators, `discriminant_switch_effect`, to the dataflow framework. While this accomplished the short-term goal of making drop elaboration more precise, it wasn't really useful in other contexts: It only supported `SwitchInt`s on the discriminant of an `enum` and did not allow effects to be applied along the "otherwise" branch. In const-propagation, for example, arbitrary edge-specific effects for the targets of a `SwitchInt` can be used to remember the value a `match` scrutinee must have in each arm.

This PR replaces `discriminant_switch_effect` with a more general `switch_int_edge_effects` method. The new method has a slightly different interface from the other edge-specific effect methods (e.g. `call_return_effect`). This divergence is explained in the new method's documentation, and reading the changes to the various dataflow impls as well as `direction.rs` should further clarify things. This PR should not change behavior.

3 files changed, 228 insertions, 132 deletions

diff --git a/compiler/rustc_mir/src/dataflow/framework/direction.rs b/compiler/rustc_mir/src/dataflow/framework/direction.rs
index 76c48100371..ca2bb6e0bf7 100644
--- a/compiler/rustc_mir/src/dataflow/framework/direction.rs
+++ b/compiler/rustc_mir/src/dataflow/framework/direction.rs
@@ -1,10 +1,10 @@
 use rustc_index::bit_set::BitSet;
 use rustc_middle::mir::{self, BasicBlock, Location};
-use rustc_middle::ty::{self, TyCtxt};
+use rustc_middle::ty::TyCtxt;
 use std::ops::RangeInclusive;
 
 use super::visitor::{ResultsVisitable, ResultsVisitor};
-use super::{Analysis, Effect, EffectIndex, GenKillAnalysis, GenKillSet};
+use super::{Analysis, Effect, EffectIndex, GenKillAnalysis, GenKillSet, SwitchIntTarget};
 
 pub trait Direction {
     fn is_forward() -> bool;
@@ -425,8 +425,8 @@ impl Direction for Forward {
 
     fn join_state_into_successors_of<A>(
         analysis: &A,
-        tcx: TyCtxt<'tcx>,
-        body: &mir::Body<'tcx>,
+        _tcx: TyCtxt<'tcx>,
+        _body: &mir::Body<'tcx>,
         dead_unwinds: Option<&BitSet<BasicBlock>>,
         exit_state: &mut A::Domain,
         (bb, bb_data): (BasicBlock, &'_ mir::BasicBlockData<'tcx>),
@@ -489,50 +489,23 @@ impl Direction for Forward {
             }
 
             SwitchInt { ref targets, ref values, ref discr, switch_ty: _ } => {
-                let enum_ = discr
-                    .place()
-                    .and_then(|discr| switch_on_enum_discriminant(tcx, &body, bb_data, discr));
-                match enum_ {
-                    // If this is a switch on an enum discriminant, a custom effect may be applied
-                    // along each outgoing edge.
-                    Some((enum_place, enum_def)) => {
-                        // 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 = analysis.bottom_value(body);
-                        let mut discriminants = enum_def.discriminants(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.clone_from(exit_state);
-                            analysis.apply_discriminant_switch_effect(
-                                &mut tmp,
-                                bb,
-                                enum_place,
-                                enum_def,
-                                variant_idx,
-                            );
-                            propagate(target, &tmp);
-                        }
-
-                        // Move out of `tmp` so we don't accidentally use it below.
-                        std::mem::drop(tmp);
-
-                        // Propagate dataflow state along the "otherwise" edge.
-                        let otherwise = targets.last().copied().unwrap();
-                        propagate(otherwise, exit_state)
-                    }
-
-                    // Otherwise, it's just a normal `SwitchInt`, and every successor sees the same
-                    // exit state.
-                    None => {
-                        for target in targets.iter().copied() {
-                            propagate(target, exit_state);
-                        }
+                let mut applier = SwitchIntEdgeEffectApplier {
+                    exit_state,
+                    targets: targets.as_ref(),
+                    values: values.as_ref(),
+                    propagate,
+                    effects_applied: false,
+                };
+
+                analysis.apply_switch_int_edge_effects(bb, discr, &mut applier);
+
+                let SwitchIntEdgeEffectApplier {
+                    exit_state, mut propagate, effects_applied, ..
+                } = applier;
+
+                if !effects_applied {
+                    for &target in targets.iter() {
+                        propagate(target, exit_state);
                     }
                 }
             }
@@ -540,37 +513,54 @@ impl Direction for Forward {
     }
 }
 
-/// Inspect a `SwitchInt`-terminated basic block to see if the condition of that `SwitchInt` is
-/// an enum discriminant.
-///
-/// We expect such blocks to have a call to `discriminant` as their last statement like so:
-///   _42 = discriminant(_1)
-///   SwitchInt(_42, ..)
-///
-/// If the basic block matches this pattern, this function returns the place corresponding to the
-/// enum (`_1` in the example above) as well as the `AdtDef` of that enum.
-fn switch_on_enum_discriminant(
-    tcx: TyCtxt<'tcx>,
-    body: &'mir mir::Body<'tcx>,
-    block: &'mir mir::BasicBlockData<'tcx>,
-    switch_on: mir::Place<'tcx>,
-) -> Option<(mir::Place<'tcx>, &'tcx ty::AdtDef)> {
-    match block.statements.last().map(|stmt| &stmt.kind) {
-        Some(mir::StatementKind::Assign(box (lhs, mir::Rvalue::Discriminant(discriminated))))
-            if *lhs == switch_on =>
-        {
-            match &discriminated.ty(body, tcx).ty.kind() {
-                ty::Adt(def, _) => Some((*discriminated, def)),
-
-                // `Rvalue::Discriminant` is also used to get the active yield point for a
-                // generator, but we do not need edge-specific effects in that case. This may
-                // change in the future.
-                ty::Generator(..) => None,
-
-                t => bug!("`discriminant` called on unexpected type {:?}", t),
-            }
+struct SwitchIntEdgeEffectApplier<'a, D, F> {
+    exit_state: &'a mut D,
+    values: &'a [u128],
+    targets: &'a [BasicBlock],
+    propagate: F,
+
+    effects_applied: bool,
+}
+
+impl<D, F> super::SwitchIntEdgeEffects<D> for SwitchIntEdgeEffectApplier<'_, D, F>
+where
+    D: Clone,
+    F: FnMut(BasicBlock, &D),
+{
+    fn apply(&mut self, mut apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget)) {
+        assert!(!self.effects_applied);
+
+        let mut tmp = None;
+        for (&value, &target) in self.values.iter().zip(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);
         }
 
-        _ => None,
+        // 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();
+        apply_edge_effect(self.exit_state, SwitchIntTarget { value: None, target: otherwise });
+        (self.propagate)(otherwise, self.exit_state);
+
+        self.effects_applied = true;
+    }
+}
+
+/// An analogue of `Option::get_or_insert_with` that stores a clone of `val` into `opt`, but uses
+/// the more efficient `clone_from` if `opt` was `Some`.
+///
+/// Returns a mutable reference to the new clone that resides in `opt`.
+//
+// FIXME: Figure out how to express this using `Option::clone_from`, or maybe lift it into the
+// standard library?
+fn opt_clone_from_or_clone<T: Clone>(opt: &'a mut Option<T>, val: &T) -> &'a mut T {
+    if opt.is_some() {
+        let ret = opt.as_mut().unwrap();
+        ret.clone_from(val);
+        ret
+    } else {
+        *opt = Some(val.clone());
+        opt.as_mut().unwrap()
     }
 }
diff --git a/compiler/rustc_mir/src/dataflow/framework/mod.rs b/compiler/rustc_mir/src/dataflow/framework/mod.rs
index eefa1395a62..65c159e6a72 100644
--- a/compiler/rustc_mir/src/dataflow/framework/mod.rs
+++ b/compiler/rustc_mir/src/dataflow/framework/mod.rs
@@ -37,8 +37,7 @@ use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::{BitSet, HybridBitSet};
 use rustc_index::vec::Idx;
 use rustc_middle::mir::{self, BasicBlock, Location};
-use rustc_middle::ty::{self, TyCtxt};
-use rustc_target::abi::VariantIdx;
+use rustc_middle::ty::TyCtxt;
 
 mod cursor;
 mod direction;
@@ -152,6 +151,8 @@ pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
     ) {
     }
 
+    /* Edge-specific effects */
+
     /// Updates the current dataflow state with the effect of a successful return from a `Call`
     /// terminator.
     ///
@@ -183,20 +184,28 @@ pub trait Analysis<'tcx>: AnalysisDomain<'tcx> {
     /// 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.
+    /// Unlike the other edge-specific effects, which are allowed to mutate `Self::Domain`
+    /// directly, overriders of this method must pass a callback to
+    /// `SwitchIntEdgeEffects::apply`. The callback will be run once for each outgoing edge and
+    /// will have access to the dataflow state that will be propagated along that edge.
+    ///
+    /// This interface is somewhat more complex than the other visitor-like "effect" methods.
+    /// However, it is both more ergonomic—callers don't need to recompute or cache information
+    /// about a given `SwitchInt` terminator for each one of its edges—and more efficient—the
+    /// engine doesn't need to clone the exit state for a block unless
+    /// `SwitchIntEdgeEffects::apply` is actually called.
     ///
     /// FIXME: This class of effects is not supported for backward dataflow analyses.
-    fn apply_discriminant_switch_effect(
+    fn apply_switch_int_edge_effects(
         &self,
-        _state: &mut Self::Domain,
         _block: BasicBlock,
-        _enum_place: mir::Place<'tcx>,
-        _adt: &ty::AdtDef,
-        _variant: VariantIdx,
+        _discr: &mir::Operand<'tcx>,
+        _apply_edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
     ) {
     }
 
+    /* Extension methods */
+
     /// 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
@@ -267,6 +276,8 @@ pub trait GenKillAnalysis<'tcx>: Analysis<'tcx> {
     ) {
     }
 
+    /* Edge-specific effects */
+
     /// See `Analysis::apply_call_return_effect`.
     fn call_return_effect(
         &self,
@@ -286,14 +297,12 @@ pub trait GenKillAnalysis<'tcx>: Analysis<'tcx> {
     ) {
     }
 
-    /// See `Analysis::apply_discriminant_switch_effect`.
-    fn discriminant_switch_effect(
+    /// See `Analysis::apply_switch_int_edge_effects`.
+    fn switch_int_edge_effects<G: GenKill<Self::Idx>>(
         &self,
-        _state: &mut impl GenKill<Self::Idx>,
         _block: BasicBlock,
-        _enum_place: mir::Place<'tcx>,
-        _adt: &ty::AdtDef,
-        _variant: VariantIdx,
+        _discr: &mir::Operand<'tcx>,
+        _edge_effects: &mut impl SwitchIntEdgeEffects<G>,
     ) {
     }
 }
@@ -339,6 +348,8 @@ where
         self.before_terminator_effect(state, terminator, location);
     }
 
+    /* Edge-specific effects */
+
     fn apply_call_return_effect(
         &self,
         state: &mut A::Domain,
@@ -359,17 +370,17 @@ where
         self.yield_resume_effect(state, resume_block, resume_place);
     }
 
-    fn apply_discriminant_switch_effect(
+    fn apply_switch_int_edge_effects(
         &self,
-        state: &mut A::Domain,
         block: BasicBlock,
-        enum_place: mir::Place<'tcx>,
-        adt: &ty::AdtDef,
-        variant: VariantIdx,
+        discr: &mir::Operand<'tcx>,
+        edge_effects: &mut impl SwitchIntEdgeEffects<A::Domain>,
     ) {
-        self.discriminant_switch_effect(state, block, enum_place, adt, variant);
+        self.switch_int_edge_effects(block, discr, edge_effects);
     }
 
+    /* Extension methods */
+
     fn into_engine(
         self,
         tcx: TyCtxt<'tcx>,
@@ -531,5 +542,17 @@ impl EffectIndex {
     }
 }
 
+pub struct SwitchIntTarget {
+    pub value: Option<u128>,
+    pub target: BasicBlock,
+}
+
+/// A type that records the edge-specific effects for a `SwitchInt` terminator.
+pub trait SwitchIntEdgeEffects<D> {
+    /// Calls `apply_edge_effect` for each outgoing edge from a `SwitchInt` terminator and
+    /// records the results.
+    fn apply(&mut self, apply_edge_effect: impl FnMut(&mut D, SwitchIntTarget));
+}
+
 #[cfg(test)]
 mod tests;
diff --git a/compiler/rustc_mir/src/dataflow/impls/mod.rs b/compiler/rustc_mir/src/dataflow/impls/mod.rs
index 1769feaf7a5..d4b9600f766 100644
--- a/compiler/rustc_mir/src/dataflow/impls/mod.rs
+++ b/compiler/rustc_mir/src/dataflow/impls/mod.rs
@@ -6,7 +6,6 @@ use rustc_index::bit_set::BitSet;
 use rustc_index::vec::Idx;
 use rustc_middle::mir::{self, Body, Location};
 use rustc_middle::ty::{self, TyCtxt};
-use rustc_target::abi::VariantIdx;
 
 use super::MoveDataParamEnv;
 
@@ -19,6 +18,7 @@ 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;
+use crate::dataflow::framework::SwitchIntEdgeEffects;
 
 mod borrowed_locals;
 pub(super) mod borrows;
@@ -352,24 +352,46 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         );
     }
 
-    fn discriminant_switch_effect(
+    fn switch_int_edge_effects<G: GenKill<Self::Idx>>(
         &self,
-        trans: &mut impl GenKill<Self::Idx>,
-        _block: mir::BasicBlock,
-        enum_place: mir::Place<'tcx>,
-        _adt: &ty::AdtDef,
-        variant: VariantIdx,
+        block: mir::BasicBlock,
+        discr: &mir::Operand<'tcx>,
+        edge_effects: &mut impl SwitchIntEdgeEffects<G>,
     ) {
-        // Kill all move paths that correspond to variants we know to be inactive along this
-        // particular outgoing edge of a `SwitchInt`.
-        drop_flag_effects::on_all_inactive_variants(
-            self.tcx,
-            self.body,
-            self.move_data(),
-            enum_place,
-            variant,
-            |mpi| trans.kill(mpi),
-        );
+        let enum_ = discr.place().and_then(|discr| {
+            switch_on_enum_discriminant(self.tcx, &self.body, &self.body[block], discr)
+        });
+
+        let (enum_place, enum_def) = match enum_ {
+            Some(x) => x,
+            None => return,
+        };
+
+        let mut discriminants = enum_def.discriminants(self.tcx);
+        edge_effects.apply(|trans, edge| {
+            let value = match edge.value {
+                Some(x) => x,
+                None => return,
+            };
+
+            // 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 (variant, _) = discriminants
+                .find(|&(_, discr)| discr.val == value)
+                .expect("Order of `AdtDef::discriminants` differed from `SwitchInt::values`");
+
+            // Kill all move paths that correspond to variants we know to be inactive along this
+            // particular outgoing edge of a `SwitchInt`.
+            drop_flag_effects::on_all_inactive_variants(
+                self.tcx,
+                self.body,
+                self.move_data(),
+                enum_place,
+                variant,
+                |mpi| trans.kill(mpi),
+            );
+        });
     }
 }
 
@@ -441,28 +463,50 @@ impl<'tcx> GenKillAnalysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         );
     }
 
-    fn discriminant_switch_effect(
+    fn switch_int_edge_effects<G: GenKill<Self::Idx>>(
         &self,
-        trans: &mut impl GenKill<Self::Idx>,
-        _block: mir::BasicBlock,
-        enum_place: mir::Place<'tcx>,
-        _adt: &ty::AdtDef,
-        variant: VariantIdx,
+        block: mir::BasicBlock,
+        discr: &mir::Operand<'tcx>,
+        edge_effects: &mut impl SwitchIntEdgeEffects<G>,
     ) {
         if !self.mark_inactive_variants_as_uninit {
             return;
         }
 
-        // Mark all move paths that correspond to variants other than this one as maybe
-        // uninitialized (in reality, they are *definitely* uninitialized).
-        drop_flag_effects::on_all_inactive_variants(
-            self.tcx,
-            self.body,
-            self.move_data(),
-            enum_place,
-            variant,
-            |mpi| trans.gen(mpi),
-        );
+        let enum_ = discr.place().and_then(|discr| {
+            switch_on_enum_discriminant(self.tcx, &self.body, &self.body[block], discr)
+        });
+
+        let (enum_place, enum_def) = match enum_ {
+            Some(x) => x,
+            None => return,
+        };
+
+        let mut discriminants = enum_def.discriminants(self.tcx);
+        edge_effects.apply(|trans, edge| {
+            let value = match edge.value {
+                Some(x) => x,
+                None => return,
+            };
+
+            // 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 (variant, _) = discriminants
+                .find(|&(_, discr)| discr.val == value)
+                .expect("Order of `AdtDef::discriminants` differed from `SwitchInt::values`");
+
+            // Mark all move paths that correspond to variants other than this one as maybe
+            // uninitialized (in reality, they are *definitely* uninitialized).
+            drop_flag_effects::on_all_inactive_variants(
+                self.tcx,
+                self.body,
+                self.move_data(),
+                enum_place,
+                variant,
+                |mpi| trans.gen(mpi),
+            );
+        });
     }
 }
 
@@ -624,3 +668,42 @@ impl<'tcx> GenKillAnalysis<'tcx> for EverInitializedPlaces<'_, 'tcx> {
         }
     }
 }
+
+/// Inspect a `SwitchInt`-terminated basic block to see if the condition of that `SwitchInt` is
+/// an enum discriminant.
+///
+/// We expect such blocks to have a call to `discriminant` as their last statement like so:
+///
+/// ```text
+/// ...
+/// _42 = discriminant(_1)
+/// SwitchInt(_42, ..)
+/// ```
+///
+/// If the basic block matches this pattern, this function returns the place corresponding to the
+/// enum (`_1` in the example above) as well as the `AdtDef` of that enum.
+fn switch_on_enum_discriminant(
+    tcx: TyCtxt<'tcx>,
+    body: &'mir mir::Body<'tcx>,
+    block: &'mir mir::BasicBlockData<'tcx>,
+    switch_on: mir::Place<'tcx>,
+) -> Option<(mir::Place<'tcx>, &'tcx ty::AdtDef)> {
+    match block.statements.last().map(|stmt| &stmt.kind) {
+        Some(mir::StatementKind::Assign(box (lhs, mir::Rvalue::Discriminant(discriminated))))
+            if *lhs == switch_on =>
+        {
+            match &discriminated.ty(body, tcx).ty.kind() {
+                ty::Adt(def, _) => Some((*discriminated, def)),
+
+                // `Rvalue::Discriminant` is also used to get the active yield point for a
+                // generator, but we do not need edge-specific effects in that case. This may
+                // change in the future.
+                ty::Generator(..) => None,
+
+                t => bug!("`discriminant` called on unexpected type {:?}", t),
+            }
+        }
+
+        _ => None,
+    }
+}