Fix unsound optimization with explicit variant discriminants

2 files changed, 32 insertions, 4 deletions

diff --git a/compiler/rustc_mir_transform/src/simplify_try.rs b/compiler/rustc_mir_transform/src/simplify_try.rs
index fd36671b36f..8da90a432ce 100644
--- a/compiler/rustc_mir_transform/src/simplify_try.rs
+++ b/compiler/rustc_mir_transform/src/simplify_try.rs
@@ -706,12 +706,22 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
         let helper = |rhs: &Rvalue<'tcx>,
                       place: &Place<'tcx>,
                       variant_index: &VariantIdx,
+                      switch_value: u128,
                       side_to_choose| {
             let place_type = place.ty(self.body, self.tcx).ty;
             let adt = match *place_type.kind() {
                 ty::Adt(adt, _) if adt.is_enum() => adt,
                 _ => return StatementEquality::NotEqual,
             };
+            let variant_discr = adt.discriminant_for_variant(self.tcx, *variant_index).val;
+            if variant_discr != switch_value {
+                trace!(
+                    "NO: variant discriminant {} does not equal switch value {}",
+                    variant_discr,
+                    switch_value
+                );
+                return StatementEquality::NotEqual;
+            }
             let variant_is_fieldless = adt.variants[*variant_index].fields.is_empty();
             if !variant_is_fieldless {
                 trace!("NO: variant {:?} was not fieldless", variant_index);
@@ -742,18 +752,18 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
                 StatementKind::SetDiscriminant { place, variant_index },
             )
             // we need to make sure that the switch value that targets the bb with SetDiscriminant (y), is the same as the variant index
-            if Some(variant_index.index() as u128) == y_target_and_value.value => {
+            if y_target_and_value.value.is_some() => {
                 // choose basic block of x, as that has the assign
-                helper(rhs, place, variant_index, x_target_and_value.target)
+                helper(rhs, place, variant_index, y_target_and_value.value.unwrap(), x_target_and_value.target)
             }
             (
                 StatementKind::SetDiscriminant { place, variant_index },
                 StatementKind::Assign(box (_, rhs)),
             )
             // we need to make sure that the switch value that targets the bb with SetDiscriminant (x), is the same as the variant index
-            if Some(variant_index.index() as u128) == x_target_and_value.value  => {
+            if x_target_and_value.value.is_some() => {
                 // choose basic block of y, as that has the assign
-                helper(rhs, place, variant_index, y_target_and_value.target)
+                helper(rhs, place, variant_index, x_target_and_value.value.unwrap(), y_target_and_value.target)
             }
             _ => {
                 trace!("NO: statements `{:?}` and `{:?}` not considered equal", x, y);
diff --git a/src/test/ui/mir/issue-89485.rs b/src/test/ui/mir/issue-89485.rs
new file mode 100644
index 00000000000..cb507eefebb
--- /dev/null
+++ b/src/test/ui/mir/issue-89485.rs
@@ -0,0 +1,18 @@
+// Regression test for issue #89485.
+
+// run-pass
+
+#[derive(Debug, Eq, PartialEq)]
+pub enum Type {
+    A = 1,
+    B = 2,
+}
+pub fn encode(v: Type) -> Type {
+    match v {
+        Type::A => Type::B,
+        _ => v,
+    }
+}
+fn main() {
+  assert_eq!(Type::B, encode(Type::A));
+}