Extract handle_set_discriminant.

2 files changed, 35 insertions, 17 deletions

diff --git a/compiler/rustc_mir_dataflow/src/value_analysis.rs b/compiler/rustc_mir_dataflow/src/value_analysis.rs
index ecdc6a6580e..7b92eb05fba 100644
--- a/compiler/rustc_mir_dataflow/src/value_analysis.rs
+++ b/compiler/rustc_mir_dataflow/src/value_analysis.rs
@@ -67,8 +67,8 @@ pub trait ValueAnalysis<'tcx> {
             StatementKind::Assign(box (place, rvalue)) => {
                 self.handle_assign(*place, rvalue, state);
             }
-            StatementKind::SetDiscriminant { box ref place, .. } => {
-                state.flood_discr(place.as_ref(), self.map());
+            StatementKind::SetDiscriminant { box place, variant_index } => {
+                self.handle_set_discriminant(*place, *variant_index, state);
             }
             StatementKind::Intrinsic(box intrinsic) => {
                 self.handle_intrinsic(intrinsic, state);
@@ -94,6 +94,24 @@ pub trait ValueAnalysis<'tcx> {
         }
     }
 
+    fn handle_set_discriminant(
+        &self,
+        place: Place<'tcx>,
+        variant_index: VariantIdx,
+        state: &mut State<Self::Value>,
+    ) {
+        self.super_set_discriminant(place, variant_index, state)
+    }
+
+    fn super_set_discriminant(
+        &self,
+        place: Place<'tcx>,
+        _variant_index: VariantIdx,
+        state: &mut State<Self::Value>,
+    ) {
+        state.flood_discr(place.as_ref(), self.map());
+    }
+
     fn handle_intrinsic(
         &self,
         intrinsic: &NonDivergingIntrinsic<'tcx>,
diff --git a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
index 5ef1bd00a6f..7adfc9dff2a 100644
--- a/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
+++ b/compiler/rustc_mir_transform/src/dataflow_const_prop.rs
@@ -79,22 +79,22 @@ impl<'tcx> ValueAnalysis<'tcx> for ConstAnalysis<'_, 'tcx> {
         &self.map
     }
 
-    fn handle_statement(&self, statement: &Statement<'tcx>, state: &mut State<Self::Value>) {
-        match statement.kind {
-            StatementKind::SetDiscriminant { box ref place, variant_index } => {
-                state.flood_discr(place.as_ref(), &self.map);
-                if self.map.find_discr(place.as_ref()).is_some() {
-                    let enum_ty = place.ty(self.local_decls, self.tcx).ty;
-                    if let Some(discr) = self.eval_discriminant(enum_ty, variant_index) {
-                        state.assign_discr(
-                            place.as_ref(),
-                            ValueOrPlace::Value(FlatSet::Elem(discr)),
-                            &self.map,
-                        );
-                    }
-                }
+    fn handle_set_discriminant(
+        &self,
+        place: Place<'tcx>,
+        variant_index: VariantIdx,
+        state: &mut State<Self::Value>,
+    ) {
+        state.flood_discr(place.as_ref(), &self.map);
+        if self.map.find_discr(place.as_ref()).is_some() {
+            let enum_ty = place.ty(self.local_decls, self.tcx).ty;
+            if let Some(discr) = self.eval_discriminant(enum_ty, variant_index) {
+                state.assign_discr(
+                    place.as_ref(),
+                    ValueOrPlace::Value(FlatSet::Elem(discr)),
+                    &self.map,
+                );
             }
-            _ => self.super_statement(statement, state),
         }
     }