Auto merge of #76844 - simonvandel:fix-76803, r=wesleywiser

Fix #76803 miscompilation

Fixes #76803
Seems like it was an oversight that the discriminant value being set was not compared to the target value from the SwitchInt, as a comment says this is a requirement for the optimization to be sound.

r? `@wesleywiser` since you are probably familiar with the optimization and made #76837 to workaround the bug

4 files changed, 96 insertions, 35 deletions

diff --git a/compiler/rustc_mir/src/transform/simplify_try.rs b/compiler/rustc_mir/src/transform/simplify_try.rs
index 9a11d927240..4935997eb82 100644
--- a/compiler/rustc_mir/src/transform/simplify_try.rs
+++ b/compiler/rustc_mir/src/transform/simplify_try.rs
@@ -16,7 +16,7 @@ use rustc_middle::mir::visit::{NonUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
 use rustc_middle::ty::{self, List, Ty, TyCtxt};
 use rustc_target::abi::VariantIdx;
-use std::iter::{Enumerate, Peekable};
+use std::iter::{once, Enumerate, Peekable};
 use std::slice::Iter;
 
 /// Simplifies arms of form `Variant(x) => Variant(x)` to just a move.
@@ -551,6 +551,12 @@ struct SimplifyBranchSameOptimization {
     bb_to_opt_terminator: BasicBlock,
 }
 
+struct SwitchTargetAndValue {
+    target: BasicBlock,
+    // None in case of the `otherwise` case
+    value: Option<u128>,
+}
+
 struct SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
     body: &'a Body<'tcx>,
     tcx: TyCtxt<'tcx>,
@@ -562,8 +568,16 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
             .basic_blocks()
             .iter_enumerated()
             .filter_map(|(bb_idx, bb)| {
-                let (discr_switched_on, targets) = match &bb.terminator().kind {
-                    TerminatorKind::SwitchInt { targets, discr, .. } => (discr, targets),
+                let (discr_switched_on, targets_and_values) = match &bb.terminator().kind {
+                    TerminatorKind::SwitchInt { targets, discr, values, .. } => {
+                        // if values.len() == targets.len() - 1, we need to include None where no value is present
+                        // such that the zip does not throw away targets. If no `otherwise` case is in targets, the zip will simply throw away the added None
+                        let values_extended = values.iter().map(|x|Some(*x)).chain(once(None));
+                        let targets_and_values:Vec<_> = targets.iter().zip(values_extended)
+                            .map(|(target, value)| SwitchTargetAndValue{target:*target, value})
+                            .collect();
+                        assert_eq!(targets.len(), targets_and_values.len());
+                        (discr, targets_and_values)},
                     _ => return None,
                 };
 
@@ -587,9 +601,9 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
                     },
                 };
 
-                let mut iter_bbs_reachable = targets
+                let mut iter_bbs_reachable = targets_and_values
                     .iter()
-                    .map(|idx| (*idx, &self.body.basic_blocks()[*idx]))
+                    .map(|target_and_value| (target_and_value, &self.body.basic_blocks()[target_and_value.target]))
                     .filter(|(_, bb)| {
                         // Reaching `unreachable` is UB so assume it doesn't happen.
                         bb.terminator().kind != TerminatorKind::Unreachable
@@ -603,16 +617,16 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
                     })
                     .peekable();
 
-                let bb_first = iter_bbs_reachable.peek().map(|(idx, _)| *idx).unwrap_or(targets[0]);
+                let bb_first = iter_bbs_reachable.peek().map(|(idx, _)| *idx).unwrap_or(&targets_and_values[0]);
                 let mut all_successors_equivalent = StatementEquality::TrivialEqual;
 
                 // All successor basic blocks must be equal or contain statements that are pairwise considered equal.
-                for ((bb_l_idx,bb_l), (bb_r_idx,bb_r)) in iter_bbs_reachable.tuple_windows() {
+                for ((target_and_value_l,bb_l), (target_and_value_r,bb_r)) in iter_bbs_reachable.tuple_windows() {
                     let trivial_checks = bb_l.is_cleanup == bb_r.is_cleanup
                     && bb_l.terminator().kind == bb_r.terminator().kind;
                     let statement_check = || {
                         bb_l.statements.iter().zip(&bb_r.statements).try_fold(StatementEquality::TrivialEqual, |acc,(l,r)| {
-                            let stmt_equality = self.statement_equality(*adt_matched_on, &l, bb_l_idx, &r, bb_r_idx, self.tcx.sess.opts.debugging_opts.mir_opt_level);
+                            let stmt_equality = self.statement_equality(*adt_matched_on, &l, target_and_value_l, &r, target_and_value_r);
                             if matches!(stmt_equality, StatementEquality::NotEqual) {
                                 // short circuit
                                 None
@@ -634,7 +648,7 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
                         // statements are trivially equal, so just take first
                         trace!("Statements are trivially equal");
                         Some(SimplifyBranchSameOptimization {
-                            bb_to_goto: bb_first,
+                            bb_to_goto: bb_first.target,
                             bb_to_opt_terminator: bb_idx,
                         })
                     }
@@ -669,10 +683,9 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
         &self,
         adt_matched_on: Place<'tcx>,
         x: &Statement<'tcx>,
-        x_bb_idx: BasicBlock,
+        x_target_and_value: &SwitchTargetAndValue,
         y: &Statement<'tcx>,
-        y_bb_idx: BasicBlock,
-        mir_opt_level: usize,
+        y_target_and_value: &SwitchTargetAndValue,
     ) -> StatementEquality {
         let helper = |rhs: &Rvalue<'tcx>,
                       place: &Place<'tcx>,
@@ -691,13 +704,7 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
 
             match rhs {
                 Rvalue::Use(operand) if operand.place() == Some(adt_matched_on) => {
-                    // FIXME(76803): This logic is currently broken because it does not take into
-                    // account the current discriminant value.
-                    if mir_opt_level > 2 {
-                        StatementEquality::ConsideredEqual(side_to_choose)
-                    } else {
-                        StatementEquality::NotEqual
-                    }
+                    StatementEquality::ConsideredEqual(side_to_choose)
                 }
                 _ => {
                     trace!(
@@ -717,16 +724,20 @@ impl<'a, 'tcx> SimplifyBranchSameOptimizationFinder<'a, 'tcx> {
             (
                 StatementKind::Assign(box (_, rhs)),
                 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 => {
                 // choose basic block of x, as that has the assign
-                helper(rhs, place, variant_index, x_bb_idx)
+                helper(rhs, place, variant_index, 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  => {
                 // choose basic block of y, as that has the assign
-                helper(rhs, place, variant_index, y_bb_idx)
+                helper(rhs, place, variant_index, y_target_and_value.target)
             }
             _ => {
                 trace!("NO: statements `{:?}` and `{:?}` not considered equal", x, y);
diff --git a/src/test/mir-opt/76803_regression.encode.SimplifyBranchSame.diff b/src/test/mir-opt/76803_regression.encode.SimplifyBranchSame.diff
new file mode 100644
index 00000000000..28b8329606c
--- /dev/null
+++ b/src/test/mir-opt/76803_regression.encode.SimplifyBranchSame.diff
@@ -0,0 +1,28 @@
+- // MIR for `encode` before SimplifyBranchSame
++ // MIR for `encode` after SimplifyBranchSame
+  
+  fn encode(_1: Type) -> Type {
+      debug v => _1;                       // in scope 0 at $DIR/76803_regression.rs:10:15: 10:16
+      let mut _0: Type;                    // return place in scope 0 at $DIR/76803_regression.rs:10:27: 10:31
+      let mut _2: isize;                   // in scope 0 at $DIR/76803_regression.rs:12:9: 12:16
+  
+      bb0: {
+          _2 = discriminant(_1);           // scope 0 at $DIR/76803_regression.rs:12:9: 12:16
+          switchInt(move _2) -> [0_isize: bb2, otherwise: bb1]; // scope 0 at $DIR/76803_regression.rs:12:9: 12:16
+      }
+  
+      bb1: {
+          _0 = move _1;                    // scope 0 at $DIR/76803_regression.rs:13:14: 13:15
+          goto -> bb3;                     // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
+      }
+  
+      bb2: {
+          discriminant(_0) = 1;            // scope 0 at $DIR/76803_regression.rs:12:20: 12:27
+          goto -> bb3;                     // scope 0 at $DIR/76803_regression.rs:11:5: 14:6
+      }
+  
+      bb3: {
+          return;                          // scope 0 at $DIR/76803_regression.rs:15:2: 15:2
+      }
+  }
+  
diff --git a/src/test/mir-opt/76803_regression.rs b/src/test/mir-opt/76803_regression.rs
new file mode 100644
index 00000000000..05dc3c97841
--- /dev/null
+++ b/src/test/mir-opt/76803_regression.rs
@@ -0,0 +1,19 @@
+// compile-flags: -Z mir-opt-level=1
+// EMIT_MIR 76803_regression.encode.SimplifyBranchSame.diff
+
+#[derive(Debug, Eq, PartialEq)]
+pub enum Type {
+    A,
+    B,
+}
+
+pub fn encode(v: Type) -> Type {
+    match v {
+        Type::A => Type::B,
+        _ => v,
+    }
+}
+
+fn main() {
+    assert_eq!(Type::B, encode(Type::A));
+}
diff --git a/src/test/mir-opt/simplify_arm.id.SimplifyBranchSame.diff b/src/test/mir-opt/simplify_arm.id.SimplifyBranchSame.diff
index 06f359da2e7..81a0e6ba0b4 100644
--- a/src/test/mir-opt/simplify_arm.id.SimplifyBranchSame.diff
+++ b/src/test/mir-opt/simplify_arm.id.SimplifyBranchSame.diff
@@ -13,24 +13,27 @@
   
       bb0: {
           _2 = discriminant(_1);           // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
-          switchInt(move _2) -> [0_isize: bb1, 1_isize: bb3, otherwise: bb2]; // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
+-         switchInt(move _2) -> [0_isize: bb1, 1_isize: bb3, otherwise: bb2]; // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
++         goto -> bb1;                     // scope 0 at $DIR/simplify-arm.rs:11:9: 11:16
       }
   
       bb1: {
-          discriminant(_0) = 0;            // scope 0 at $DIR/simplify-arm.rs:12:17: 12:21
-          goto -> bb4;                     // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
-      }
-  
-      bb2: {
-          unreachable;                     // scope 0 at $DIR/simplify-arm.rs:10:11: 10:12
-      }
-  
-      bb3: {
+-         discriminant(_0) = 0;            // scope 0 at $DIR/simplify-arm.rs:12:17: 12:21
+-         goto -> bb4;                     // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
+-     }
+- 
+-     bb2: {
+-         unreachable;                     // scope 0 at $DIR/simplify-arm.rs:10:11: 10:12
+-     }
+- 
+-     bb3: {
           _0 = move _1;                    // scope 1 at $DIR/simplify-arm.rs:11:20: 11:27
-          goto -> bb4;                     // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
+-         goto -> bb4;                     // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
++         goto -> bb2;                     // scope 0 at $DIR/simplify-arm.rs:10:5: 13:6
       }
   
-      bb4: {
+-     bb4: {
++     bb2: {
           return;                          // scope 0 at $DIR/simplify-arm.rs:14:2: 14:2
       }
   }