Tell LLVM about impossible niche tags

3 files changed, 465 insertions, 19 deletions

diff --git a/compiler/rustc_codegen_ssa/src/mir/operand.rs b/compiler/rustc_codegen_ssa/src/mir/operand.rs
index 7e355b6406a..fca0695f296 100644
--- a/compiler/rustc_codegen_ssa/src/mir/operand.rs
+++ b/compiler/rustc_codegen_ssa/src/mir/operand.rs
@@ -9,6 +9,7 @@ use rustc_middle::mir::{self, ConstValue};
 use rustc_middle::ty::Ty;
 use rustc_middle::ty::layout::{LayoutOf, TyAndLayout};
 use rustc_middle::{bug, span_bug};
+use rustc_session::config::OptLevel;
 use tracing::{debug, instrument};
 
 use super::place::{PlaceRef, PlaceValue};
@@ -496,6 +497,18 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                     _ => (tag_imm, bx.cx().immediate_backend_type(tag_op.layout)),
                 };
 
+                // Layout ensures that we only get here for cases where the discriminant
+                // value and the variant index match, since that's all `Niche` can encode.
+                // But for emphasis and debugging, let's double-check one anyway.
+                debug_assert_eq!(
+                    self.layout
+                        .ty
+                        .discriminant_for_variant(bx.tcx(), untagged_variant)
+                        .unwrap()
+                        .val,
+                    u128::from(untagged_variant.as_u32()),
+                );
+
                 let relative_max = niche_variants.end().as_u32() - niche_variants.start().as_u32();
 
                 // We have a subrange `niche_start..=niche_end` inside `range`.
@@ -537,6 +550,21 @@ impl<'a, 'tcx, V: CodegenObject> OperandRef<'tcx, V> {
                         relative_discr,
                         bx.cx().const_uint(tag_llty, relative_max as u64),
                     );
+
+                    // Thanks to parameter attributes and load metadata, LLVM already knows
+                    // the general valid range of the tag. It's possible, though, for there
+                    // to be an impossible value *in the middle*, which those ranges don't
+                    // communicate, so it's worth an `assume` to let the optimizer know.
+                    if niche_variants.contains(&untagged_variant)
+                        && bx.cx().sess().opts.optimize != OptLevel::No
+                    {
+                        let impossible =
+                            u64::from(untagged_variant.as_u32() - niche_variants.start().as_u32());
+                        let impossible = bx.cx().const_uint(tag_llty, impossible);
+                        let ne = bx.icmp(IntPredicate::IntNE, relative_discr, impossible);
+                        bx.assume(ne);
+                    }
+
                     (is_niche, cast_tag, niche_variants.start().as_u32() as u128)
                 };
 
diff --git a/tests/codegen/enum/enum-match.rs b/tests/codegen/enum/enum-match.rs
index a24b98050d2..bea0d5f33b0 100644
--- a/tests/codegen/enum/enum-match.rs
+++ b/tests/codegen/enum/enum-match.rs
@@ -1,7 +1,8 @@
 //@ compile-flags: -Copt-level=1
-//@ only-x86_64
+//@ only-64bit
 
 #![crate_type = "lib"]
+#![feature(core_intrinsics)]
 
 // Check each of the 3 cases for `codegen_get_discr`.
 
@@ -11,11 +12,12 @@ pub enum Enum0 {
     B,
 }
 
-// CHECK: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match0{{.*}}
+// CHECK-LABEL: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match0(i8{{.+}}%0)
 // CHECK-NEXT: start:
-// CHECK-NEXT: %1 = icmp eq i8 %0, 2
-// CHECK-NEXT: %2 = and i8 %0, 1
-// CHECK-NEXT: %{{.+}} = select i1 %1, i8 13, i8 %2
+// CHECK-NEXT: %[[IS_B:.+]] = icmp eq i8 %0, 2
+// CHECK-NEXT: %[[TRUNC:.+]] = and i8 %0, 1
+// CHECK-NEXT: %[[R:.+]] = select i1 %[[IS_B]], i8 13, i8 %[[TRUNC]]
+// CHECK-NEXT: ret i8 %[[R]]
 #[no_mangle]
 pub fn match0(e: Enum0) -> u8 {
     use Enum0::*;
@@ -32,13 +34,14 @@ pub enum Enum1 {
     C,
 }
 
-// CHECK: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match1{{.*}}
+// CHECK-LABEL: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match1(i8{{.+}}%0)
 // CHECK-NEXT: start:
-// CHECK-NEXT: %1 = add{{( nsw)?}} i8 %0, -2
-// CHECK-NEXT: %2 = zext i8 %1 to i64
-// CHECK-NEXT: %3 = icmp ult i8 %1, 2
-// CHECK-NEXT: %4 = add nuw nsw i64 %2, 1
-// CHECK-NEXT: %_2 = select i1 %3, i64 %4, i64 0
+// CHECK-NEXT: %[[REL_VAR:.+]] = add{{( nsw)?}} i8 %0, -2
+// CHECK-NEXT: %[[REL_VAR_WIDE:.+]] = zext i8 %[[REL_VAR]] to i64
+// CHECK-NEXT: %[[IS_NICHE:.+]] = icmp ult i8 %[[REL_VAR]], 2
+// CHECK-NEXT: %[[NICHE_DISCR:.+]] = add nuw nsw i64 %[[REL_VAR_WIDE]], 1
+// CHECK-NEXT: %[[DISCR:.+]] = select i1 %[[IS_NICHE]], i64 %[[NICHE_DISCR]], i64 0
+// CHECK-NEXT: switch i64 %[[DISCR]]
 #[no_mangle]
 pub fn match1(e: Enum1) -> u8 {
     use Enum1::*;
@@ -92,14 +95,14 @@ pub enum Enum2 {
     E,
 }
 
-// CHECK: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match2{{.*}}
+// CHECK-LABEL: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match2(i8{{.+}}%0)
 // CHECK-NEXT: start:
-// CHECK-NEXT: %1 = add i8 %0, 2
-// CHECK-NEXT: %2 = zext i8 %1 to i64
-// CHECK-NEXT: %3 = icmp ult i8 %1, 4
-// CHECK-NEXT: %4 = add nuw nsw i64 %2, 1
-// CHECK-NEXT: %_2 = select i1 %3, i64 %4, i64 0
-// CHECK-NEXT: switch i64 %_2, label {{.*}} [
+// CHECK-NEXT: %[[REL_VAR:.+]] = add i8 %0, 2
+// CHECK-NEXT: %[[REL_VAR_WIDE:.+]] = zext i8 %[[REL_VAR]] to i64
+// CHECK-NEXT: %[[IS_NICHE:.+]] = icmp ult i8 %[[REL_VAR]], 4
+// CHECK-NEXT: %[[NICHE_DISCR:.+]] = add nuw nsw i64 %[[REL_VAR_WIDE]], 1
+// CHECK-NEXT: %[[DISCR:.+]] = select i1 %[[IS_NICHE]], i64 %[[NICHE_DISCR]], i64 0
+// CHECK-NEXT: switch i64 %[[DISCR]]
 #[no_mangle]
 pub fn match2(e: Enum2) -> u8 {
     use Enum2::*;
@@ -111,3 +114,357 @@ pub fn match2(e: Enum2) -> u8 {
         E => 250,
     }
 }
+
+// And make sure it works even if the niched scalar is a pointer.
+// (For example, that we don't try to `sub` on pointers.)
+
+// CHECK-LABEL: define noundef{{( range\(i16 -?[0-9]+, -?[0-9]+\))?}} i16 @match3(ptr{{.+}}%0)
+// CHECK-NEXT: start:
+// CHECK-NEXT: %[[IS_NULL:.+]] = icmp eq ptr %0, null
+// CHECK-NEXT: br i1 %[[IS_NULL]]
+#[no_mangle]
+pub fn match3(e: Option<&u8>) -> i16 {
+    match e {
+        Some(r) => *r as _,
+        None => -1,
+    }
+}
+
+// If the untagged variant is in the middle, there's an impossible value that's
+// not reflected in the `range` parameter attribute, so we assume it away.
+
+#[derive(PartialEq)]
+pub enum MiddleNiche {
+    A,
+    B,
+    C(bool),
+    D,
+    E,
+}
+
+// CHECK-LABEL: define noundef{{( range\(i8 -?[0-9]+, -?[0-9]+\))?}} i8 @match4(i8{{.+}}%0)
+// CHECK-NEXT: start:
+// CHECK-NEXT: %[[REL_VAR:.+]] = add nsw i8 %0, -2
+// CHECK-NEXT: %[[IS_NICHE:.+]] = icmp ult i8 %[[REL_VAR]], 5
+// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i8 %[[REL_VAR]], 2
+// CHECK-NEXT: call void @llvm.assume(i1 %[[NOT_IMPOSSIBLE]])
+// CHECK-NEXT: %[[DISCR:.+]] = select i1 %[[IS_NICHE]], i8 %[[REL_VAR]], i8 2
+// CHECK-NEXT: switch i8 %[[DISCR]]
+#[no_mangle]
+pub fn match4(e: MiddleNiche) -> u8 {
+    use MiddleNiche::*;
+    match e {
+        A => 13,
+        B => 100,
+        C(b) => b as u8,
+        D => 200,
+        E => 250,
+    }
+}
+
+// CHECK-LABEL: define{{.+}}i1 @match4_is_c(i8{{.+}}%e)
+// CHECK-NEXT: start
+// CHECK-NEXT: %[[REL_VAR:.+]] = add nsw i8 %e, -2
+// CHECK-NEXT: %[[NOT_NICHE:.+]] = icmp ugt i8 %[[REL_VAR]], 4
+// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i8 %[[REL_VAR]], 2
+// CHECK-NEXT: call void @llvm.assume(i1 %[[NOT_IMPOSSIBLE]])
+// CHECK-NEXT: ret i1 %[[NOT_NICHE]]
+#[no_mangle]
+pub fn match4_is_c(e: MiddleNiche) -> bool {
+    // Before #139098, this couldn't optimize out the `select` because it looked
+    // like it was possible for a `2` to be produced on both sides.
+
+    std::intrinsics::discriminant_value(&e) == 2
+}
+
+// You have to do something pretty obnoxious to get a variant index that doesn't
+// fit in the tag size, but it's possible
+
+pub enum Never {}
+
+pub enum HugeVariantIndex {
+    V000(Never),
+    V001(Never),
+    V002(Never),
+    V003(Never),
+    V004(Never),
+    V005(Never),
+    V006(Never),
+    V007(Never),
+    V008(Never),
+    V009(Never),
+    V010(Never),
+    V011(Never),
+    V012(Never),
+    V013(Never),
+    V014(Never),
+    V015(Never),
+    V016(Never),
+    V017(Never),
+    V018(Never),
+    V019(Never),
+    V020(Never),
+    V021(Never),
+    V022(Never),
+    V023(Never),
+    V024(Never),
+    V025(Never),
+    V026(Never),
+    V027(Never),
+    V028(Never),
+    V029(Never),
+    V030(Never),
+    V031(Never),
+    V032(Never),
+    V033(Never),
+    V034(Never),
+    V035(Never),
+    V036(Never),
+    V037(Never),
+    V038(Never),
+    V039(Never),
+    V040(Never),
+    V041(Never),
+    V042(Never),
+    V043(Never),
+    V044(Never),
+    V045(Never),
+    V046(Never),
+    V047(Never),
+    V048(Never),
+    V049(Never),
+    V050(Never),
+    V051(Never),
+    V052(Never),
+    V053(Never),
+    V054(Never),
+    V055(Never),
+    V056(Never),
+    V057(Never),
+    V058(Never),
+    V059(Never),
+    V060(Never),
+    V061(Never),
+    V062(Never),
+    V063(Never),
+    V064(Never),
+    V065(Never),
+    V066(Never),
+    V067(Never),
+    V068(Never),
+    V069(Never),
+    V070(Never),
+    V071(Never),
+    V072(Never),
+    V073(Never),
+    V074(Never),
+    V075(Never),
+    V076(Never),
+    V077(Never),
+    V078(Never),
+    V079(Never),
+    V080(Never),
+    V081(Never),
+    V082(Never),
+    V083(Never),
+    V084(Never),
+    V085(Never),
+    V086(Never),
+    V087(Never),
+    V088(Never),
+    V089(Never),
+    V090(Never),
+    V091(Never),
+    V092(Never),
+    V093(Never),
+    V094(Never),
+    V095(Never),
+    V096(Never),
+    V097(Never),
+    V098(Never),
+    V099(Never),
+    V100(Never),
+    V101(Never),
+    V102(Never),
+    V103(Never),
+    V104(Never),
+    V105(Never),
+    V106(Never),
+    V107(Never),
+    V108(Never),
+    V109(Never),
+    V110(Never),
+    V111(Never),
+    V112(Never),
+    V113(Never),
+    V114(Never),
+    V115(Never),
+    V116(Never),
+    V117(Never),
+    V118(Never),
+    V119(Never),
+    V120(Never),
+    V121(Never),
+    V122(Never),
+    V123(Never),
+    V124(Never),
+    V125(Never),
+    V126(Never),
+    V127(Never),
+    V128(Never),
+    V129(Never),
+    V130(Never),
+    V131(Never),
+    V132(Never),
+    V133(Never),
+    V134(Never),
+    V135(Never),
+    V136(Never),
+    V137(Never),
+    V138(Never),
+    V139(Never),
+    V140(Never),
+    V141(Never),
+    V142(Never),
+    V143(Never),
+    V144(Never),
+    V145(Never),
+    V146(Never),
+    V147(Never),
+    V148(Never),
+    V149(Never),
+    V150(Never),
+    V151(Never),
+    V152(Never),
+    V153(Never),
+    V154(Never),
+    V155(Never),
+    V156(Never),
+    V157(Never),
+    V158(Never),
+    V159(Never),
+    V160(Never),
+    V161(Never),
+    V162(Never),
+    V163(Never),
+    V164(Never),
+    V165(Never),
+    V166(Never),
+    V167(Never),
+    V168(Never),
+    V169(Never),
+    V170(Never),
+    V171(Never),
+    V172(Never),
+    V173(Never),
+    V174(Never),
+    V175(Never),
+    V176(Never),
+    V177(Never),
+    V178(Never),
+    V179(Never),
+    V180(Never),
+    V181(Never),
+    V182(Never),
+    V183(Never),
+    V184(Never),
+    V185(Never),
+    V186(Never),
+    V187(Never),
+    V188(Never),
+    V189(Never),
+    V190(Never),
+    V191(Never),
+    V192(Never),
+    V193(Never),
+    V194(Never),
+    V195(Never),
+    V196(Never),
+    V197(Never),
+    V198(Never),
+    V199(Never),
+    V200(Never),
+    V201(Never),
+    V202(Never),
+    V203(Never),
+    V204(Never),
+    V205(Never),
+    V206(Never),
+    V207(Never),
+    V208(Never),
+    V209(Never),
+    V210(Never),
+    V211(Never),
+    V212(Never),
+    V213(Never),
+    V214(Never),
+    V215(Never),
+    V216(Never),
+    V217(Never),
+    V218(Never),
+    V219(Never),
+    V220(Never),
+    V221(Never),
+    V222(Never),
+    V223(Never),
+    V224(Never),
+    V225(Never),
+    V226(Never),
+    V227(Never),
+    V228(Never),
+    V229(Never),
+    V230(Never),
+    V231(Never),
+    V232(Never),
+    V233(Never),
+    V234(Never),
+    V235(Never),
+    V236(Never),
+    V237(Never),
+    V238(Never),
+    V239(Never),
+    V240(Never),
+    V241(Never),
+    V242(Never),
+    V243(Never),
+    V244(Never),
+    V245(Never),
+    V246(Never),
+    V247(Never),
+    V248(Never),
+    V249(Never),
+    V250(Never),
+    V251(Never),
+    V252(Never),
+    V253(Never),
+    V254(Never),
+    V255(Never),
+    V256(Never),
+
+    Possible257,
+    Bool258(bool),
+    Possible259,
+}
+
+// CHECK-LABEL: define noundef{{( range\(i8 [0-9]+, [0-9]+\))?}} i8 @match5(i8{{.+}}%0)
+// CHECK-NEXT: start:
+// CHECK-NEXT: %[[REL_VAR:.+]] = add nsw i8 %0, -2
+// CHECK-NEXT: %[[REL_VAR_WIDE:.+]] = zext i8 %[[REL_VAR]] to i64
+// CHECK-NEXT: %[[IS_NICHE:.+]] = icmp ult i8 %[[REL_VAR]], 3
+// CHECK-NEXT: %[[NOT_IMPOSSIBLE:.+]] = icmp ne i8 %[[REL_VAR]], 1
+// CHECK-NEXT: call void @llvm.assume(i1 %[[NOT_IMPOSSIBLE]])
+// CHECK-NEXT: %[[NICHE_DISCR:.+]] = add nuw nsw i64 %[[REL_VAR_WIDE]], 257
+// CHECK-NEXT: %[[DISCR:.+]] = select i1 %[[IS_NICHE]], i64 %[[NICHE_DISCR]], i64 258
+// CHECK-NEXT: switch i64 %[[DISCR]],
+// CHECK-NEXT:   i64 257,
+// CHECK-NEXT:   i64 258,
+// CHECK-NEXT:   i64 259,
+#[no_mangle]
+pub fn match5(e: HugeVariantIndex) -> u8 {
+    use HugeVariantIndex::*;
+    match e {
+        Possible257 => 13,
+        Bool258(b) => b as u8,
+        Possible259 => 100,
+    }
+}
diff --git a/tests/codegen/enum/enum-two-variants-match.rs b/tests/codegen/enum/enum-two-variants-match.rs
index 21ae1f96bca..5b6c0e2f8ad 100644
--- a/tests/codegen/enum/enum-two-variants-match.rs
+++ b/tests/codegen/enum/enum-two-variants-match.rs
@@ -1,5 +1,5 @@
 //@ compile-flags: -Copt-level=3 -C no-prepopulate-passes
-//@ only-x86_64 (because these discriminants are isize)
+//@ only-64bit (because these discriminants are isize)
 
 #![crate_type = "lib"]
 
@@ -51,3 +51,64 @@ pub fn result_match(x: Result<u64, i64>) -> u16 {
         Ok(_) => 42,
     }
 }
+
+// CHECK-LABEL: @option_bool_match(
+#[no_mangle]
+pub fn option_bool_match(x: Option<bool>) -> char {
+    // CHECK: %[[RAW:.+]] = load i8, ptr %x
+    // CHECK: %[[IS_NONE:.+]] = icmp eq i8 %[[RAW]], 2
+    // CHECK: %[[OPT_DISCR:.+]] = select i1 %[[IS_NONE]], i64 0, i64 1
+    // CHECK: %[[OPT_DISCR_T:.+]] = trunc nuw i64 %[[OPT_DISCR]] to i1
+    // CHECK: br i1 %[[OPT_DISCR_T]], label %[[BB_SOME:.+]], label %[[BB_NONE:.+]]
+
+    // CHECK: [[BB_SOME]]:
+    // CHECK: %[[FIELD:.+]] = load i8, ptr %x
+    // CHECK: %[[FIELD_T:.+]] = trunc nuw i8 %[[FIELD]] to i1
+    // CHECK: br i1 %[[FIELD_T]]
+    match x {
+        None => 'n',
+        Some(false) => 'f',
+        Some(true) => 't',
+    }
+}
+
+use std::cmp::Ordering::{self, *};
+// CHECK-LABEL: @option_ordering_match(
+#[no_mangle]
+pub fn option_ordering_match(x: Option<Ordering>) -> char {
+    // CHECK: %[[RAW:.+]] = load i8, ptr %x
+    // CHECK: %[[IS_NONE:.+]] = icmp eq i8 %[[RAW]], 2
+    // CHECK: %[[OPT_DISCR:.+]] = select i1 %[[IS_NONE]], i64 0, i64 1
+    // CHECK: %[[OPT_DISCR_T:.+]] = trunc nuw i64 %[[OPT_DISCR]] to i1
+    // CHECK: br i1 %[[OPT_DISCR_T]], label %[[BB_SOME:.+]], label %[[BB_NONE:.+]]
+
+    // CHECK: [[BB_SOME]]:
+    // CHECK: %[[FIELD:.+]] = load i8, ptr %x
+    // CHECK: switch i8 %[[FIELD]], label %[[UNREACHABLE:.+]] [
+    // CHECK-NEXT: i8 -1, label
+    // CHECK-NEXT: i8 0, label
+    // CHECK-NEXT: i8 1, label
+    // CHECK-NEXT: ]
+
+    // CHECK: [[UNREACHABLE]]:
+    // CHECK-NEXT: unreachable
+    match x {
+        None => '?',
+        Some(Less) => '<',
+        Some(Equal) => '=',
+        Some(Greater) => '>',
+    }
+}
+
+// CHECK-LABEL: @option_nonzero_match(
+#[no_mangle]
+pub fn option_nonzero_match(x: Option<std::num::NonZero<u16>>) -> u16 {
+    // CHECK: %[[IS_NONE:.+]] = icmp eq i16 %x, 0
+    // CHECK: %[[OPT_DISCR:.+]] = select i1 %[[IS_NONE]], i64 0, i64 1
+    // CHECK: %[[OPT_DISCR_T:.+]] = trunc nuw i64 %[[OPT_DISCR]] to i1
+    // CHECK: br i1 %[[OPT_DISCR_T]], label %[[BB_SOME:.+]], label %[[BB_NONE:.+]]
+    match x {
+        None => 123,
+        Some(_) => 987,
+    }
+}