Simplify dataflow `SwitchInt` handling.

Current `SwitchInt` handling has complicated control flow.

- The dataflow engine calls `Analysis::apply_switch_int_edge_effects`,
  passing in an "applier" that impls `SwitchIntEdgeEffects`.
- `apply_switch_int_edge_effects` possibly calls `apply` on the applier,
  passing it a closure.
- The `apply` method calls the closure on each `SwitchInt` edge.
- The closure operates on the edge.

I.e. control flow goes from the engine, to the analysis, to the applier
(which came from the engine), to the closure (which came from the
analysis). It took me a while to work this out.

This commit changes to a simpler structure that maintains the important
characteristics.

- The dataflow engine calls `Analysis::get_switch_int_data`.
- `get_switch_int_data` returns an `Option<Self::SwitchIntData>` value.
- If that returned value was `Some`, the dataflow engine calls
  `Analysis::apply_switch_int_edge_effect` on each edge, passing the
  `Self::SwitchIntData` value.
- `Analysis::apply_switch_int_edge_effect` operates on the edge.

I.e. control flow goes from the engine, to the analysis, to the
engine, to the analysis.

Added:
- The `Analysis::SwitchIntData` assoc type and the
  `Analysis::get_switch_int_data` method. Both only need to be
  defined by analyses that look at `SwitchInt` terminators.
- The `MaybePlacesSwitchIntData` struct, which has three fields.

Changes:
- `Analysis::apply_switch_int_edge_effects` becomes
  `Analysis::apply_switch_int_edge_effect`, which is a little simpler
  because it's dealing with a single edge instead of all edges.

Removed:
- The `SwitchIntEdgeEffects` trait, and its two impls:
  `BackwardSwitchIntEdgeEffectsApplier` (which has six fields) and
  `ForwardSwitchIntEdgeEffectsApplier` structs (which has four fields).
- The closure.

The new structure is more concise and simpler.

1 files changed, 79 insertions, 58 deletions

diff --git a/compiler/rustc_mir_dataflow/src/impls/initialized.rs b/compiler/rustc_mir_dataflow/src/impls/initialized.rs
index 9fdbbb1578e..d5982b9cba2 100644
--- a/compiler/rustc_mir_dataflow/src/impls/initialized.rs
+++ b/compiler/rustc_mir_dataflow/src/impls/initialized.rs
@@ -1,20 +1,45 @@
 use std::assert_matches::assert_matches;
 
+use rustc_abi::VariantIdx;
 use rustc_index::Idx;
 use rustc_index::bit_set::{BitSet, MixedBitSet};
 use rustc_middle::bug;
 use rustc_middle::mir::{self, Body, CallReturnPlaces, Location, TerminatorEdges};
+use rustc_middle::ty::util::Discr;
 use rustc_middle::ty::{self, TyCtxt};
 use tracing::{debug, instrument};
 
 use crate::elaborate_drops::DropFlagState;
-use crate::framework::SwitchIntEdgeEffects;
+use crate::framework::SwitchIntTarget;
 use crate::move_paths::{HasMoveData, InitIndex, InitKind, LookupResult, MoveData, MovePathIndex};
 use crate::{
     Analysis, GenKill, MaybeReachable, drop_flag_effects, drop_flag_effects_for_function_entry,
     drop_flag_effects_for_location, on_all_children_bits, on_lookup_result_bits,
 };
 
+// Used by both `MaybeInitializedPlaces` and `MaybeUninitializedPlaces`.
+pub struct MaybePlacesSwitchIntData<'tcx> {
+    enum_place: mir::Place<'tcx>,
+    discriminants: Vec<(VariantIdx, Discr<'tcx>)>,
+    index: usize,
+}
+
+impl<'tcx> MaybePlacesSwitchIntData<'tcx> {
+    // The discriminant order in the `SwitchInt` targets should match the order yielded by
+    // `AdtDef::discriminants`. We rely on this to match each discriminant in the targets to its
+    // corresponding variant in linear time.
+    fn next_discr(&mut self, value: u128) -> VariantIdx {
+        // An out-of-bounds abort will occur if the discriminant ordering isn't as described above.
+        loop {
+            let (variant, discr) = self.discriminants[self.index];
+            self.index += 1;
+            if discr.val == value {
+                return variant;
+            }
+        }
+    }
+}
+
 /// `MaybeInitializedPlaces` tracks all places that might be
 /// initialized upon reaching a particular point in the control flow
 /// for a function.
@@ -247,6 +272,8 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
     /// We use a mixed bitset to avoid paying too high a memory footprint.
     type Domain = MaybeReachable<MixedBitSet<MovePathIndex>>;
 
+    type SwitchIntData = MaybePlacesSwitchIntData<'tcx>;
+
     const NAME: &'static str = "maybe_init";
 
     fn bottom_value(&self, _: &mir::Body<'tcx>) -> Self::Domain {
@@ -328,46 +355,42 @@ impl<'tcx> Analysis<'tcx> for MaybeInitializedPlaces<'_, 'tcx> {
         });
     }
 
-    fn apply_switch_int_edge_effects(
+    fn get_switch_int_data(
         &mut self,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
-        edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
-    ) {
+    ) -> Option<Self::SwitchIntData> {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
-            return;
+            return None;
         }
 
-        let enum_ = discr.place().and_then(|discr| {
-            switch_on_enum_discriminant(self.tcx, self.body, &self.body[block], discr)
-        });
-
-        let Some((enum_place, enum_def)) = enum_ else {
-            return;
-        };
-
-        let mut discriminants = enum_def.discriminants(self.tcx);
-        edge_effects.apply(|state, edge| {
-            let Some(value) = edge.value else {
-                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`");
+        discr.place().and_then(|discr| {
+            switch_on_enum_discriminant(self.tcx, self.body, &self.body[block], discr).map(
+                |(enum_place, enum_def)| MaybePlacesSwitchIntData {
+                    enum_place,
+                    discriminants: enum_def.discriminants(self.tcx).collect(),
+                    index: 0,
+                },
+            )
+        })
+    }
 
+    fn apply_switch_int_edge_effect(
+        &mut self,
+        data: &mut Self::SwitchIntData,
+        state: &mut Self::Domain,
+        edge: SwitchIntTarget,
+    ) {
+        if let Some(value) = edge.value {
             // 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.move_data(),
-                enum_place,
-                variant,
+                self.move_data,
+                data.enum_place,
+                data.next_discr(value),
                 |mpi| state.kill(mpi),
             );
-        });
+        }
     }
 }
 
@@ -378,6 +401,8 @@ pub type MaybeUninitializedPlacesDomain = MixedBitSet<MovePathIndex>;
 impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
     type Domain = MaybeUninitializedPlacesDomain;
 
+    type SwitchIntData = MaybePlacesSwitchIntData<'tcx>;
+
     const NAME: &'static str = "maybe_uninit";
 
     fn bottom_value(&self, _: &mir::Body<'tcx>) -> Self::Domain {
@@ -447,50 +472,46 @@ impl<'tcx> Analysis<'tcx> for MaybeUninitializedPlaces<'_, 'tcx> {
         });
     }
 
-    fn apply_switch_int_edge_effects(
+    fn get_switch_int_data(
         &mut self,
         block: mir::BasicBlock,
         discr: &mir::Operand<'tcx>,
-        edge_effects: &mut impl SwitchIntEdgeEffects<Self::Domain>,
-    ) {
+    ) -> Option<Self::SwitchIntData> {
         if !self.tcx.sess.opts.unstable_opts.precise_enum_drop_elaboration {
-            return;
+            return None;
         }
 
         if !self.mark_inactive_variants_as_uninit {
-            return;
+            return None;
         }
 
-        let enum_ = discr.place().and_then(|discr| {
-            switch_on_enum_discriminant(self.tcx, self.body, &self.body[block], discr)
-        });
-
-        let Some((enum_place, enum_def)) = enum_ else {
-            return;
-        };
-
-        let mut discriminants = enum_def.discriminants(self.tcx);
-        edge_effects.apply(|state, edge| {
-            let Some(value) = edge.value else {
-                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`");
+        discr.place().and_then(|discr| {
+            switch_on_enum_discriminant(self.tcx, self.body, &self.body[block], discr).map(
+                |(enum_place, enum_def)| MaybePlacesSwitchIntData {
+                    enum_place,
+                    discriminants: enum_def.discriminants(self.tcx).collect(),
+                    index: 0,
+                },
+            )
+        })
+    }
 
+    fn apply_switch_int_edge_effect(
+        &mut self,
+        data: &mut Self::SwitchIntData,
+        state: &mut Self::Domain,
+        edge: SwitchIntTarget,
+    ) {
+        if let Some(value) = edge.value {
             // 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.move_data(),
-                enum_place,
-                variant,
+                self.move_data,
+                data.enum_place,
+                data.next_discr(value),
                 |mpi| state.gen_(mpi),
             );
-        });
+        }
     }
 }