Add a macro that derives `TryFrom<u32>` for fieldless enums

2 files changed, 63 insertions, 0 deletions

diff --git a/compiler/rustc_macros/src/lib.rs b/compiler/rustc_macros/src/lib.rs
index f46c795b956..0df674eb4c9 100644
--- a/compiler/rustc_macros/src/lib.rs
+++ b/compiler/rustc_macros/src/lib.rs
@@ -20,6 +20,7 @@ mod lift;
 mod query;
 mod serialize;
 mod symbols;
+mod try_from;
 mod type_foldable;
 mod type_visitable;
 
@@ -165,3 +166,12 @@ decl_derive!(
         suggestion_part,
         applicability)] => diagnostics::subdiagnostic_derive
 );
+
+decl_derive! {
+    [TryFromU32] =>
+    /// Derives `TryFrom<u32>` for the annotated `enum`, which must have no fields.
+    /// Each variant maps to the value it would produce under an `as u32` cast.
+    ///
+    /// The error type is `u32`.
+    try_from::try_from_u32
+}
diff --git a/compiler/rustc_macros/src/try_from.rs b/compiler/rustc_macros/src/try_from.rs
new file mode 100644
index 00000000000..fde7c8f6204
--- /dev/null
+++ b/compiler/rustc_macros/src/try_from.rs
@@ -0,0 +1,53 @@
+use proc_macro2::TokenStream;
+use quote::{quote, quote_spanned};
+use syn::Data;
+use syn::spanned::Spanned;
+use synstructure::Structure;
+
+pub(crate) fn try_from_u32(s: Structure<'_>) -> TokenStream {
+    let span_error = |span, message: &str| {
+        quote_spanned! { span => const _: () = ::core::compile_error!(#message); }
+    };
+
+    // Must be applied to an enum type.
+    if let Some(span) = match &s.ast().data {
+        Data::Enum(_) => None,
+        Data::Struct(s) => Some(s.struct_token.span()),
+        Data::Union(u) => Some(u.union_token.span()),
+    } {
+        return span_error(span, "type is not an enum (TryFromU32)");
+    }
+
+    // The enum's variants must not have fields.
+    let variant_field_errors = s
+        .variants()
+        .iter()
+        .filter_map(|v| v.ast().fields.iter().map(|f| f.span()).next())
+        .map(|span| span_error(span, "enum variant cannot have fields (TryFromU32)"))
+        .collect::<TokenStream>();
+    if !variant_field_errors.is_empty() {
+        return variant_field_errors;
+    }
+
+    let ctor = s
+        .variants()
+        .iter()
+        .map(|v| v.construct(|_, _| -> TokenStream { unreachable!() }))
+        .collect::<Vec<_>>();
+    s.gen_impl(quote! {
+        // The surrounding code might have shadowed these identifiers.
+        use ::core::convert::TryFrom;
+        use ::core::primitive::u32;
+        use ::core::result::Result::{self, Ok, Err};
+
+        gen impl TryFrom<u32> for @Self {
+            type Error = u32;
+
+            #[allow(deprecated)] // Don't warn about deprecated variants.
+            fn try_from(value: u32) -> Result<Self, Self::Error> {
+                #( if value == const { #ctor as u32 } { return Ok(#ctor) } )*
+                Err(value)
+            }
+        }
+    })
+}