SimplifyArmIdentity only for locals with the same type

Co-Authored-By: Mazdak Farrokhzad <twingoow@gmail.com>

Author: tmiasko

src/librustc_mir/transform/simplify_try.rs

@@ -34,7 +34,8 @@ pub struct SimplifyArmIdentity;
 
 impl<'tcx> MirPass<'tcx> for SimplifyArmIdentity {
     fn run_pass(&self, _: TyCtxt<'tcx>, _: MirSource<'tcx>, body: &mut Body<'tcx>) {
-        for bb in body.basic_blocks_mut() {
+        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+        for bb in basic_blocks {
             // Need 3 statements:
             let (s0, s1, s2) = match &mut *bb.statements {
                 [s0, s1, s2] => (s0, s1, s2),
@@ -51,7 +52,12 @@ impl<'tcx> MirPass<'tcx> for SimplifyArmIdentity {
                 Some(x) => x,
             };
             if local_tmp_s0 != local_tmp_s1
+                // The field-and-variant information match up.
                 || vf_s0 != vf_s1
+                // Source and target locals have the same type.
+                // FIXME(Centril | oli-obk): possibly relax to same layout?
+                || local_decls[local_0].ty != local_decls[local_1].ty
+                // We're setting the discriminant of `local_0` to this variant.
                 || Some((local_0, vf_s0.var_idx)) != match_set_discr(s2)
             {
                 continue;

src/test/mir-opt/simplify-arm-identity.rs

@@ -0,0 +1,75 @@
+// Checks that `SimplifyArmIdentity` is not applied if enums have incompatible layouts.
+// Regression test for issue #66856.
+//
+// compile-flags: -Zmir-opt-level=2
+
+enum Src {
+    Foo(u8),
+    Bar,
+}
+
+enum Dst {
+    Foo(u8),
+}
+
+fn main() {
+    let e: Src = Src::Foo(0);
+    let _: Dst = match e {
+        Src::Foo(x) => Dst::Foo(x),
+        Src::Bar => Dst::Foo(0),
+    };
+}
+
+// END RUST SOURCE
+// START rustc.main.SimplifyArmIdentity.before.mir
+// fn main() -> () {
+//     ...
+//     bb0: {
+//         StorageLive(_1);
+//         ((_1 as Foo).0: u8) = const 0u8;
+//         discriminant(_1) = 0;
+//         StorageLive(_2);
+//         _3 = discriminant(_1);
+//         switchInt(move _3) -> [0isize: bb3, 1isize: bb1, otherwise: bb2];
+//     }
+//     bb1: {
+//         ((_2 as Foo).0: u8) = const 0u8;
+//         discriminant(_2) = 0;
+//         goto -> bb4;
+//     }
+//     ...
+//     bb3: {
+//         _4 = ((_1 as Foo).0: u8);
+//         ((_2 as Foo).0: u8) = move _4;
+//         discriminant(_2) = 0;
+//         goto -> bb4;
+//     }
+//     ...
+// }
+// END rustc.main.SimplifyArmIdentity.before.mir
+// START rustc.main.SimplifyArmIdentity.after.mir
+// fn main() -> () {
+//     ...
+//     bb0: {
+//         StorageLive(_1);
+//         ((_1 as Foo).0: u8) = const 0u8;
+//         discriminant(_1) = 0;
+//         StorageLive(_2);
+//         _3 = discriminant(_1);
+//         switchInt(move _3) -> [0isize: bb3, 1isize: bb1, otherwise: bb2];
+//     }
+//     bb1: {
+//         ((_2 as Foo).0: u8) = const 0u8;
+//         discriminant(_2) = 0;
+//         goto -> bb4;
+//     }
+//     ...
+//     bb3: {
+//         _4 = ((_1 as Foo).0: u8);
+//         ((_2 as Foo).0: u8) = move _4;
+//         discriminant(_2) = 0;
+//         goto -> bb4;
+//     }
+//     ...
+// }
+// END rustc.main.SimplifyArmIdentity.after.mir

src/test/ui/issues/issue-66851.rs

@@ -0,0 +1,20 @@
+// This used to mis-compile because the mir-opt `SimplifyArmIdentity`
+// did not check that the types matched up in the `Ok(r)` branch.
+//
+// run-pass
+// compile-flags: -Zmir-opt-level=2
+
+#[derive(Debug, PartialEq, Eq)]
+enum SpecialsRes { Res(u64) }
+
+fn e103() -> SpecialsRes {
+    if let Ok(r) = "1".parse() {
+        SpecialsRes::Res(r)
+    } else {
+        SpecialsRes::Res(42)
+    }
+}
+
+fn main() {
+    assert_eq!(e103(), SpecialsRes::Res(1));
+}