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

code review
1

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()));
    };
}
2

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
3

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
4

Fixed:

1 Like
5

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.