Attribute macro to derive repr for enum and implement discriminant method for it

After this issue I whipped up a proc macro to derive a repr and appropriate discriminant getter method (defaults to discriminant) for an arbitrary enum.

//! Attribute macro to implement a discriminant method for enums with a specific representation type.

use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::{ToTokens, quote};
use syn::parse::{Parse, ParseStream};
use syn::token::Comma;
use syn::{Data, DeriveInput, Ident, Type, parse_macro_input};

const DEFAULT_DISCRIMINANT_METHOD_NAME: &str = "discriminant";

/// Attribute macro to implement a discriminant method for enums with a specific representation type.
///
/// # Panics
///
/// This macro will panic if the input type is not an enum or if the arguments are not specified correctly.
#[proc_macro_attribute]
pub fn repr_discriminant(args: TokenStream, input: TokenStream) -> TokenStream {
    let args: Args = parse_macro_input!(args);
    let typ = args.typ;
    let method_name = args
        .method_name
        .unwrap_or_else(|| Ident::new(DEFAULT_DISCRIMINANT_METHOD_NAME, Span::call_site()));
    let input: DeriveInput = parse_macro_input!(input as DeriveInput);
    let name = &input.ident;
    let body = match input.data {
        Data::Enum(_) => {
            quote! {
                pub const fn #method_name(&self) -> #typ {
                    // SAFETY: The macro guarantees that the enum is repr(#typ).
                    unsafe {
                        *::core::ptr::from_ref(self)
                            .cast::<#typ>()
                    }
                }
            }
        }
        _ => unimplemented!(),
    };

    let mut tokens = quote! {
        #[repr(#typ)]
    };
    tokens.extend(input.to_token_stream());
    tokens.extend(quote! {
        impl #name {
            #body
        }
    });
    TokenStream::from(tokens)
}

/// Arguments for the `repr_discriminant` macro.
struct Args {
    typ: Type,
    method_name: Option<Ident>,
}

impl Parse for Args {
    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
        let typ: Type = input.parse()?;

        if input.is_empty() {
            return Ok(Self {
                typ,
                method_name: None,
            });
        };

        let _: Comma = input.parse()?;

        Ok(Self {
            typ,
            method_name: Some(input.parse()?),
        })
    }
}

Is the macro okay?
Did I reinvent the wheel?

Usage

use repr_derive::repr_discriminant;

#[repr_discriminant(u8, id)]
enum Foo {
    Bar(u32) = 0xab,
}

fn main() {
    let foo = Foo::Bar(42);
    println!("{:#04X}", foo.id());
}
~/macro_test> cargo expand
    Checking macro_test v0.1.0 (/home/neumann/macro_test)
    Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.08s

#![feature(prelude_import)]
#[prelude_import]
use std::prelude::rust_2024::*;
#[macro_use]
extern crate std;
use repr_derive::repr_discriminant;
#[repr(u8)]
enum Foo {
    Bar(u32) = 0xab,
}
impl Foo {
    pub const fn id(&self) -> u8 {
        unsafe { *::core::ptr::from_ref(self).cast::<u8>() }
    }
}
fn main() {
    let foo = Foo::Bar(42);
    {
        ::std::io::_print(format_args!("{0:#04X}\n", foo.id()));
    };
}

Can't input be interpolated as well? Might easier to read than extending the token stream:

    let input = input.to_token_stream();

    let tokens = quote! {
        #[repr(#typ)]
        #input

        impl #name {
            #body
        }
    };

The macro also doesn't consider generic parameters of the enum. I.e. I can't slap it onto a

enum Foo<'a> {
    Bar(&'a str) = 0xab,
}

for example.

1 Like

why not use strum EnumDiscriminants?

you can convert an Enum to its EnumDiscriminant value and then use the as u8 on the discriminant.


use strum::{EnumDiscriminants};

#[derive(
    EnumDiscriminants)]
#[repr(u8)]
enum Test {
    Value1(f64) = 5,
    Value2 = 6
}

fn main() {
    
    let my_enum = Test::Value1(25.0);
    let discriminant : TestDiscriminants = my_enum.into();
    let my_id = discriminant as u8;

    println!("{my_id}");
}
1 Like

Fixed:

1 Like

This generates a whole new type and thus is overkill for my use case.
Additionally, it required me to perform an extra step in order to get the integer value of the discriminant.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.