repr_with_fallback/

lib.rs

//! Automatically generate [`From`] and [`Into`] impls for `enum`s with custom discriminant values
//! and a fallback variant.
//!
//! # Usage
//! ```
//! use repr_with_fallback::repr_with_fallback;
//!
//! repr_with_fallback! {
//!     /// A DNSSEC algorithm.
//!     #[derive(Debug, PartialEq)]
//!     pub enum Algorithm {
//!         /// ...
//!         RSASHA256 = 8,
//!         RSASHA512 = 10,
//!         ECDSAP256SHA256 = 13,
//!         ECDSAP384SHA384 = 14,
//!         ED25519 = 15,
//!         Unassigned(u8),
//!     }
//! }
//!
//! assert_eq!(u8::from(Algorithm::ED25519), 15);
//! assert_eq!(Algorithm::from(15), Algorithm::ED25519);
//!
//! assert_eq!(u8::from(Algorithm::Unassigned(17)), 17);
//! assert_eq!(Algorithm::from(17), Algorithm::Unassigned(17));
//! ```
//!
//! There are two restrictions imposed on the `enum`:
//! 1. It must have only unit variants, except for exactly one variant with exactly one unnamed
//!    field. This is used as the fallback variant, and the type of its field must match the type of
//!    the discriminants.
//! 2. Every variant must have a discriminant value provided.
//!
//! The repr type does not need to be numerical:
//! ```
//! repr_with_fallback! {
//!     pub enum Strings {
//!         Foo = "static",
//!         Bar = "string",
//!         Baz = "slices",
//!         Spam = "work",
//!         Eggs = "too",
//!         Unknown(&'static str),
//!     }
//! }
//!
//! let s: &'static str = Strings::Foo.into();
//! assert_eq!(s, "static");
//! ```

use proc_macro::TokenStream;
use quote::{quote, quote_spanned};
use syn::{Expr, Fields, FieldsUnnamed, ItemEnum, Type, Variant};

#[proc_macro]
pub fn repr_with_fallback(input: TokenStream) -> TokenStream {
    let ast: ItemEnum = match syn::parse(input) {
        Ok(ast) => ast,
        Err(_) => {
            return quote! {
                compile_error!("This macro expects an enum definition as its input.");
            }
            .into();
        }
    };

    let (fallback_variant, repr_type) = match get_repr_type(&ast) {
        Ok(r) => r,
        Err(e) => return e.into(),
    };

    let (unit_variants, discriminant_exprs) = match get_discriminant_exprs(&ast) {
        Ok(d) => d,
        Err(e) => return e.into(),
    };

    // create new enum definition, this time without the discriminants (they would cause a compile
    // error)
    let mut enum_without_discriminants = ast.clone();
    enum_without_discriminants
        .variants
        .iter_mut()
        .for_each(|var| var.discriminant = None);

    let from_enum_impl = gen_from_enum_impl(
        &ast,
        repr_type,
        fallback_variant,
        &unit_variants,
        &discriminant_exprs,
    );
    let from_repr_impl = gen_from_repr_impl(
        &ast,
        repr_type,
        fallback_variant,
        &unit_variants,
        &discriminant_exprs,
    );

    quote! {
        #enum_without_discriminants
        #from_enum_impl
        #from_repr_impl
    }
    .into()
}

fn get_repr_type(ast: &ItemEnum) -> Result<(&Variant, &Type), proc_macro2::TokenStream> {
    let variants = &ast.variants;
    let unit_variants_count = variants
        .iter()
        .filter(|var| var.fields == Fields::Unit)
        .count();

    let err = quote_spanned! {ast.ident.span()=>
        compile_error!("Tthe enum must have only unit variants plus exactly one variant with exactly one unnamed field.");
    };

    if unit_variants_count != variants.len() - 1 {
        return Err(err);
    }

    // check that we have one fallback variant
    let (fallback_variant, fallback_variant_fields) = variants
        .iter()
        .filter_map(|var| match &var.fields {
            Fields::Unnamed(FieldsUnnamed { unnamed, .. }) => Some((var, unnamed)),
            _ => None,
        })
        .next()
        .ok_or_else(|| err.clone())?;
    // check that the fallback variant has one field
    if fallback_variant_fields.len() != 1 {
        return Err(err);
    }

    Ok((fallback_variant, &fallback_variant_fields[0].ty))
}

fn get_discriminant_exprs(
    ast: &ItemEnum,
) -> Result<(Vec<&Variant>, Vec<&Expr>), proc_macro2::TokenStream> {
    let unit_variants: Vec<_> = ast
        .variants
        .iter()
        .filter(|var| var.fields == Fields::Unit)
        .collect();

    // check that all unit variants have a discriminant.
    // this is either `Ok(d)` with `d` being all discriminants, or `Err(i`) with `i` being the ident
    // of the first variant without a discriminant
    let discriminants: Result<Vec<_>, _> = unit_variants
        .iter()
        .map(|var| var.discriminant.as_ref().ok_or(&var.ident))
        .collect();
    let discriminant_exprs = match discriminants {
        Err(ident) => {
            return Err(quote_spanned! {ident.span()=>
                compile_error!("All unit variants must have a discriminant.");
            })
        }
        // get the expression of the discriminant (discard the " = " part)
        Ok(d) => d.iter().map(|d| &d.1).collect(),
    };
    Ok((unit_variants, discriminant_exprs))
}

fn gen_from_enum_impl(
    ast: &ItemEnum,
    repr_type: &Type,
    fallback_variant: &Variant,
    unit_variants: &[&Variant],
    discriminant_exprs: &[&Expr],
) -> proc_macro2::TokenStream {
    let enum_ident = &ast.ident;
    let unit_variant_maps =
        unit_variants
            .iter()
            .zip(discriminant_exprs.iter())
            .map(|(var, expr)| {
                let var_ident = &var.ident;
                quote! {
                    #enum_ident::#var_ident => #expr
                }
            });
    let unit_variant_maps = quote!(#(#unit_variant_maps),*);

    let fallback_ident = &fallback_variant.ident;
    quote! {
        impl From<#enum_ident> for #repr_type {
            fn from(val: #enum_ident) -> Self {
                match val {
                    #unit_variant_maps,
                    #enum_ident::#fallback_ident(x) => x,
                }
            }
        }
    }
}

fn gen_from_repr_impl(
    ast: &ItemEnum,
    repr_type: &Type,
    fallback_variant: &Variant,
    unit_variants: &[&Variant],
    discriminant_exprs: &[&Expr],
) -> proc_macro2::TokenStream {
    let enum_ident = &ast.ident;
    let unit_variant_map_iter =
        unit_variants
            .iter()
            .zip(discriminant_exprs.iter())
            .map(|(var, expr)| {
                let var_ident = &var.ident;
                quote! {
                    #expr => #enum_ident::#var_ident
                }
            });
    let unit_variant_maps = quote!(#(#unit_variant_map_iter),*);

    let fallback_ident = &fallback_variant.ident;
    quote! {
        impl From<#repr_type> for #enum_ident {
            fn from(val: #repr_type) -> Self {
                match val {
                    #unit_variant_maps,
                    x => #enum_ident::#fallback_ident(x),
                }
            }
        }
    }
}