Rollup merge of #73418 - doctorn:variants-intrinsic, r=kennytm

Add unstable `core::mem::variant_count` intrinsic

Adds a new `const fn` intrinsic which can be used to determine the number of variants in an `enum`.

I've shown this to a couple of people and they invariably ask 'why on earth?', but there's actually a very neat use case:

At the moment, if you want to create an opaque array type that's indexed by an `enum` with one element for each variant, you either have to hard-code the number of variants, add a `LENGTH` variant or use a `Vec`, none of which are suitable in general (number of variants could change; pattern matching `LENGTH` becomes frustrating; might not have `alloc`). By including this intrinsic, it becomes possible to write the following:

```rust
#[derive(Copy, Clone)]
enum OpaqueIndex {
    A = 0,
    B,
    C,
}

struct OpaqueVec<T>(Box<[T; std::mem::num_variants::<OpaqueIndex>()]>);

impl<T> std::ops::Index<OpaqueIndex> for OpaqueVec<T> {
    type Output = T;

    fn index(&self, idx: OpaqueIndex) -> &Self::Output {
        &self.0[idx as usize]
    }
}
```

(We even have a use cases for this in `rustc` and I plan to use it to re-implement the lang-items table.)

8 files changed, 111 insertions, 5 deletions

diff --git a/src/libcore/intrinsics.rs b/src/libcore/intrinsics.rs
index 2298958b881..5d090187591 100644
--- a/src/libcore/intrinsics.rs
+++ b/src/libcore/intrinsics.rs
@@ -1917,6 +1917,15 @@ extern "rust-intrinsic" {
     #[rustc_const_unstable(feature = "const_discriminant", issue = "69821")]
     pub fn discriminant_value<T>(v: &T) -> <T as DiscriminantKind>::Discriminant;
 
+    /// Returns the number of variants of the type `T` cast to a `usize`;
+    /// if `T` has no variants, returns 0. Uninhabited variants will be counted.
+    ///
+    /// The to-be-stabilized version of this intrinsic is
+    /// [`std::mem::variant_count`](../../std/mem/fn.variant_count.html)
+    #[rustc_const_unstable(feature = "variant_count", issue = "73662")]
+    #[cfg(not(bootstrap))]
+    pub fn variant_count<T>() -> usize;
+
     /// Rust's "try catch" construct which invokes the function pointer `try_fn`
     /// with the data pointer `data`.
     ///
@@ -1960,6 +1969,12 @@ extern "rust-intrinsic" {
     pub fn ptr_guaranteed_ne<T>(ptr: *const T, other: *const T) -> bool;
 }
 
+#[rustc_const_unstable(feature = "variant_count", issue = "73662")]
+#[cfg(bootstrap)]
+pub const fn variant_count<T>() -> usize {
+    0
+}
+
 // Some functions are defined here because they accidentally got made
 // available in this module on stable. See <https://github.com/rust-lang/rust/issues/15702>.
 // (`transmute` also falls into this category, but it cannot be wrapped due to the
diff --git a/src/libcore/lib.rs b/src/libcore/lib.rs
index 4eb2fdbd078..2b26e5303a8 100644
--- a/src/libcore/lib.rs
+++ b/src/libcore/lib.rs
@@ -124,6 +124,7 @@
 #![feature(unsized_locals)]
 #![feature(untagged_unions)]
 #![feature(unwind_attributes)]
+#![feature(variant_count)]
 #![feature(doc_alias)]
 #![feature(mmx_target_feature)]
 #![feature(tbm_target_feature)]
diff --git a/src/libcore/mem/mod.rs b/src/libcore/mem/mod.rs
index 066bb8b3dc7..1bd7ae3a34e 100644
--- a/src/libcore/mem/mod.rs
+++ b/src/libcore/mem/mod.rs
@@ -999,3 +999,33 @@ impl<T> fmt::Debug for Discriminant<T> {
 pub const fn discriminant<T>(v: &T) -> Discriminant<T> {
     Discriminant(intrinsics::discriminant_value(v))
 }
+
+/// Returns the number of variants in the enum type `T`.
+///
+/// If `T` is not an enum, calling this function will not result in undefined behavior, but the
+/// return value is unspecified. Equally, if `T` is an enum with more variants than `usize::MAX`
+/// the return value is unspecified. Uninhabited variants will be counted.
+///
+/// # Examples
+///
+/// ```
+/// # #![feature(never_type)]
+/// # #![feature(variant_count)]
+///
+/// use std::mem;
+///
+/// enum Void {}
+/// enum Foo { A(&'static str), B(i32), C(i32) }
+///
+/// assert_eq!(mem::variant_count::<Void>(), 0);
+/// assert_eq!(mem::variant_count::<Foo>(), 3);
+///
+/// assert_eq!(mem::variant_count::<Option<!>>(), 2);
+/// assert_eq!(mem::variant_count::<Result<!, !>>(), 2);
+/// ```
+#[inline(always)]
+#[unstable(feature = "variant_count", issue = "73662")]
+#[rustc_const_unstable(feature = "variant_count", issue = "73662")]
+pub const fn variant_count<T>() -> usize {
+    intrinsics::variant_count::<T>()
+}
diff --git a/src/librustc_codegen_llvm/intrinsic.rs b/src/librustc_codegen_llvm/intrinsic.rs
index 0a1cc31044a..130c0cf1877 100644
--- a/src/librustc_codegen_llvm/intrinsic.rs
+++ b/src/librustc_codegen_llvm/intrinsic.rs
@@ -206,7 +206,7 @@ impl IntrinsicCallMethods<'tcx> for Builder<'a, 'll, 'tcx> {
                 }
             }
             "size_of" | "pref_align_of" | "min_align_of" | "needs_drop" | "type_id"
-            | "type_name" => {
+            | "type_name" | "variant_count" => {
                 let value = self
                     .tcx
                     .const_eval_instance(ty::ParamEnv::reveal_all(), instance, None)
diff --git a/src/librustc_mir/interpret/intrinsics.rs b/src/librustc_mir/interpret/intrinsics.rs
index 6ac1e6be036..88ba28dab82 100644
--- a/src/librustc_mir/interpret/intrinsics.rs
+++ b/src/librustc_mir/interpret/intrinsics.rs
@@ -69,6 +69,13 @@ crate fn eval_nullary_intrinsic<'tcx>(
             ConstValue::from_machine_usize(n, &tcx)
         }
         sym::type_id => ConstValue::from_u64(tcx.type_id_hash(tp_ty)),
+        sym::variant_count => {
+            if let ty::Adt(ref adt, _) = tp_ty.kind {
+                ConstValue::from_machine_usize(adt.variants.len() as u64, &tcx)
+            } else {
+                ConstValue::from_machine_usize(0u64, &tcx)
+            }
+        }
         other => bug!("`{}` is not a zero arg intrinsic", other),
     })
 }
@@ -109,10 +116,13 @@ impl<'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> InterpCx<'mir, 'tcx, M> {
             | sym::needs_drop
             | sym::size_of
             | sym::type_id
-            | sym::type_name => {
+            | sym::type_name
+            | sym::variant_count => {
                 let gid = GlobalId { instance, promoted: None };
                 let ty = match intrinsic_name {
-                    sym::min_align_of | sym::pref_align_of | sym::size_of => self.tcx.types.usize,
+                    sym::min_align_of | sym::pref_align_of | sym::size_of | sym::variant_count => {
+                        self.tcx.types.usize
+                    }
                     sym::needs_drop => self.tcx.types.bool,
                     sym::type_id => self.tcx.types.u64,
                     sym::type_name => self.tcx.mk_static_str(),
diff --git a/src/librustc_span/symbol.rs b/src/librustc_span/symbol.rs
index fa1368b104c..857734037af 100644
--- a/src/librustc_span/symbol.rs
+++ b/src/librustc_span/symbol.rs
@@ -832,6 +832,7 @@ symbols! {
         v1,
         val,
         var,
+        variant_count,
         vec,
         Vec,
         version,
diff --git a/src/librustc_typeck/check/intrinsic.rs b/src/librustc_typeck/check/intrinsic.rs
index ef6c7c14404..1c0b22ca737 100644
--- a/src/librustc_typeck/check/intrinsic.rs
+++ b/src/librustc_typeck/check/intrinsic.rs
@@ -75,7 +75,7 @@ pub fn intrinsic_operation_unsafety(intrinsic: &str) -> hir::Unsafety {
         | "saturating_sub" | "rotate_left" | "rotate_right" | "ctpop" | "ctlz" | "cttz"
         | "bswap" | "bitreverse" | "discriminant_value" | "type_id" | "likely" | "unlikely"
         | "ptr_guaranteed_eq" | "ptr_guaranteed_ne" | "minnumf32" | "minnumf64" | "maxnumf32"
-        | "maxnumf64" | "type_name" => hir::Unsafety::Normal,
+        | "maxnumf64" | "type_name" | "variant_count" => hir::Unsafety::Normal,
         _ => hir::Unsafety::Unsafe,
     }
 }
@@ -137,7 +137,9 @@ pub fn check_intrinsic_type(tcx: TyCtxt<'_>, it: &hir::ForeignItem<'_>) {
         let unsafety = intrinsic_operation_unsafety(&name[..]);
         let (n_tps, inputs, output) = match &name[..] {
             "breakpoint" => (0, Vec::new(), tcx.mk_unit()),
-            "size_of" | "pref_align_of" | "min_align_of" => (1, Vec::new(), tcx.types.usize),
+            "size_of" | "pref_align_of" | "min_align_of" | "variant_count" => {
+                (1, Vec::new(), tcx.types.usize)
+            }
             "size_of_val" | "min_align_of_val" => {
                 (1, vec![tcx.mk_imm_ptr(param(0))], tcx.types.usize)
             }
diff --git a/src/test/ui/consts/const-variant-count.rs b/src/test/ui/consts/const-variant-count.rs
new file mode 100644
index 00000000000..455419d2c7f
--- /dev/null
+++ b/src/test/ui/consts/const-variant-count.rs
@@ -0,0 +1,47 @@
+// run-pass
+#![allow(dead_code)]
+#![feature(variant_count)]
+#![feature(never_type)]
+
+use std::mem::variant_count;
+
+enum Void {}
+
+enum Foo {
+    A,
+    B,
+    C,
+}
+
+enum Bar {
+    A,
+    B,
+    C,
+    D(usize),
+    E { field_1: usize, field_2: Foo },
+}
+
+struct Baz {
+    a: u32,
+    b: *const u8,
+}
+
+const TEST_VOID: usize = variant_count::<Void>();
+const TEST_FOO: usize = variant_count::<Foo>();
+const TEST_BAR: usize = variant_count::<Bar>();
+
+const NO_ICE_STRUCT: usize = variant_count::<Baz>();
+const NO_ICE_BOOL: usize = variant_count::<bool>();
+const NO_ICE_PRIM: usize = variant_count::<*const u8>();
+
+fn main() {
+    assert_eq!(TEST_VOID, 0);
+    assert_eq!(TEST_FOO, 3);
+    assert_eq!(TEST_BAR, 5);
+    assert_eq!(variant_count::<Void>(), 0);
+    assert_eq!(variant_count::<Foo>(), 3);
+    assert_eq!(variant_count::<Bar>(), 5);
+    assert_eq!(variant_count::<Option<char>>(), 2);
+    assert_eq!(variant_count::<Option<!>>(), 2);
+    assert_eq!(variant_count::<Result<!, !>>(), 2);
+}