Make the enum check work for negative discriminants

The discriminant check was not working correctly for negative numbers.
This change fixes that by masking out the relevant bits correctly.

Author: 1c3t3a

compiler/rustc_mir_transform/src/check_enums.rs

@@ -120,6 +120,7 @@ enum EnumCheckType<'tcx> {
     },
 }
 
+#[derive(Debug, Copy, Clone)]
 struct TyAndSize<'tcx> {
     pub ty: Ty<'tcx>,
     pub size: Size,
@@ -338,7 +339,7 @@ fn insert_direct_enum_check<'tcx>(
     let invalid_discr_block_data = BasicBlockData::new(None, false);
     let invalid_discr_block = basic_blocks.push(invalid_discr_block_data);
     let block_data = &mut basic_blocks[current_block];
-    let discr = insert_discr_cast_to_u128(
+    let discr_place = insert_discr_cast_to_u128(
         tcx,
         local_decls,
         block_data,
@@ -349,13 +350,41 @@ fn insert_direct_enum_check<'tcx>(
         source_info,
     );
 
+    // Mask out the bits of the discriminant type.
+    let mask = match discr.size.bytes() {
+        1 => u8::MAX as u128,
+        2 => u16::MAX as u128,
+        4 => u32::MAX as u128,
+        8 => u64::MAX as u128,
+        16 => u128::MAX as u128,
+        invalid => bug!("Found discriminant with invalid size, has {} bytes", invalid),
+    };
+
+    let discr_masked =
+        local_decls.push(LocalDecl::with_source_info(tcx.types.u128, source_info)).into();
+    let rvalue = Rvalue::BinaryOp(
+        BinOp::BitAnd,
+        Box::new((
+            Operand::Copy(discr_place),
+            Operand::Constant(Box::new(ConstOperand {
+                span: source_info.span,
+                user_ty: None,
+                const_: Const::Val(ConstValue::from_u128(mask), tcx.types.u128),
+            })),
+        )),
+    );
+    block_data.statements.push(Statement {
+        source_info,
+        kind: StatementKind::Assign(Box::new((discr_masked, rvalue))),
+    });
+
     // Branch based on the discriminant value.
     block_data.terminator = Some(Terminator {
         source_info,
         kind: TerminatorKind::SwitchInt {
-            discr: Operand::Copy(discr),
+            discr: Operand::Copy(discr_masked),
             targets: SwitchTargets::new(
-                discriminants.into_iter().map(|discr| (discr, new_block)),
+                discriminants.into_iter().map(|discr| (discr & mask, new_block)),
                 invalid_discr_block,
             ),
         },
@@ -372,7 +401,7 @@ fn insert_direct_enum_check<'tcx>(
             })),
             expected: true,
             target: new_block,
-            msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Copy(discr))),
+            msg: Box::new(AssertKind::InvalidEnumConstruction(Operand::Copy(discr_masked))),
             // This calls panic_invalid_enum_construction, which is #[rustc_nounwind].
             // We never want to insert an unwind into unsafe code, because unwinding could
             // make a failing UB check turn into much worse UB when we start unwinding.

tests/ui/mir/enum/negative_discr_ok.rs

@@ -0,0 +1,15 @@
+//@ run-pass
+//@ compile-flags: -C debug-assertions
+
+#[allow(dead_code)]
+enum Foo {
+    A = -2,
+    B = 1,
+}
+
+fn main() {
+    let _val: Foo =
+        unsafe { std::mem::transmute::<i8, Foo>(-2) };
+    let _val: Foo =
+        unsafe { std::mem::transmute::<i8, Foo>(1) };
+}