bitflags2-derive 0.2.2

use proc_macro2::TokenStream as TokenStream2;
use quote::{format_ident, quote};

use crate::parse::FlagsInput;

/// Generates the public flag type and all supported operations.
pub(crate) fn generate(input: FlagsInput) -> TokenStream2 {
    let vis = input.vis;
    let enum_ident = input.ident;
    let variants = input.variants;
    let forwarded_attrs = input.forwarded_attrs;
    let forwarded_derives = input.forwarded_derives;

    let derive_attr = derive_attr(forwarded_derives);

    let backing = match input.explicit_backing {
        Some(ident) => quote! { #ident },
        None => {
            let max_value = variants
                .iter()
                .filter_map(|variant| variant.value_literal)
                .max()
                .unwrap_or(0);
            backing_type(max_value)
        }
    };

    let consts = variants.iter().filter(|v| !v.is_ignored()).map(|variant| {
        let ident = &variant.ident;
        let value_tokens = &variant.value_tokens;
        quote! {
            #[allow(non_upper_case_globals)]
            pub const #ident: Self = Self((#value_tokens) as #backing);
        }
    });

    let debug_arms = variants
        .iter()
        .filter(|variant| !variant.is_ignored() && variant.value_literal != Some(0))
        .map(|variant| {
            let ident = &variant.ident;
            let name = variant.ident.to_string();
            quote! {
                if self.has_flag(Self::#ident) {
                    if !first {
                        f.write_str(" | ")?;
                    }
                    f.write_str(#name)?;
                    first = false;
                }
            }
        });

    let zero_name = variants
        .iter()
        .find(|variant| variant.value_literal == Some(0))
        .map(|variant| variant.ident.to_string())
        .unwrap_or_else(|| "0".to_string());

    let all_terms = variants.iter().filter(|v| !v.is_ignored()).map(|v| {
        let ident = &v.ident;
        quote! { Self::#ident.0 }
    });

    let impl_ints_macro = format_ident!("__bitflags2_impl_ints_for_{}", enum_ident);

    quote! {
        #(#forwarded_attrs)*
        #derive_attr
        #vis struct #enum_ident(#backing);

        impl #enum_ident {
            #(#consts)*

            pub const fn empty() -> Self {
                Self(0)
            }

            pub const fn all() -> Self {
                Self(#(#all_terms)|*)
            }

            pub const fn bits(self) -> #backing {
                self.0
            }

            pub const fn from_bits(bits: #backing) -> Self {
                Self(bits)
            }

            pub const fn has_flag(self, other: Self) -> bool {
                (self.0 & other.0) == other.0
            }
        }

        macro_rules! #impl_ints_macro {
            ($($ty:ty),*) => {
                $(
                    impl ::core::convert::From<$ty> for #enum_ident {
                        fn from(bits: $ty) -> Self {
                            Self(bits as #backing)
                        }
                    }

                    impl ::core::convert::From<#enum_ident> for $ty {
                        fn from(flags: #enum_ident) -> Self {
                            flags.0 as $ty
                        }
                    }

                    impl ::core::cmp::PartialEq<$ty> for #enum_ident {
                        fn eq(&self, rhs: &$ty) -> bool {
                            self.0 as u128 == *rhs as u128
                        }
                    }

                    impl ::core::cmp::PartialEq<#enum_ident> for $ty {
                        fn eq(&self, rhs: &#enum_ident) -> bool {
                            *self as u128 == rhs.0 as u128
                        }
                    }
                )*
            };
        }

        #impl_ints_macro! {
            u8, u16, u32, u64, u128
        }

        impl ::core::ops::BitOr for #enum_ident {
            type Output = Self;

            fn bitor(self, rhs: Self) -> Self::Output {
                Self(self.0 | rhs.0)
            }
        }

        impl ::core::ops::BitOrAssign for #enum_ident {
            fn bitor_assign(&mut self, rhs: Self) {
                self.0 |= rhs.0;
            }
        }

        impl ::core::ops::BitAnd for #enum_ident {
            type Output = Self;

            fn bitand(self, rhs: Self) -> Self::Output {
                Self(self.0 & rhs.0)
            }
        }

        impl ::core::ops::BitAndAssign for #enum_ident {
            fn bitand_assign(&mut self, rhs: Self) {
                self.0 &= rhs.0;
            }
        }

        impl ::core::ops::BitXor for #enum_ident {
            type Output = Self;

            fn bitxor(self, rhs: Self) -> Self::Output {
                Self(self.0 ^ rhs.0)
            }
        }

        impl ::core::ops::BitXorAssign for #enum_ident {
            fn bitxor_assign(&mut self, rhs: Self) {
                self.0 ^= rhs.0;
            }
        }

        impl ::core::ops::Not for #enum_ident {
            type Output = Self;

            fn not(self) -> Self::Output {
                Self(!self.0)
            }
        }

        impl ::core::fmt::Debug for #enum_ident {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                let mut first = true;
                #(#debug_arms)*
                if first {
                    f.write_str(#zero_name)?;
                }
                Ok(())
            }
        }
    }
}

fn derive_attr(forwarded_derives: Vec<syn::Path>) -> TokenStream2 {
    if forwarded_derives.is_empty() {
        quote! { #[derive(Copy, Clone, PartialEq, Eq)] }
    } else {
        quote! { #[derive(Copy, Clone, PartialEq, Eq, #(#forwarded_derives),*)] }
    }
}

fn backing_type(max_value: u128) -> TokenStream2 {
    if max_value <= u8::MAX as u128 {
        quote! { u8 }
    } else if max_value <= u16::MAX as u128 {
        quote! { u16 }
    } else if max_value <= u32::MAX as u128 {
        quote! { u32 }
    } else if max_value <= u64::MAX as u128 {
        quote! { u64 }
    } else {
        quote! { u128 }
    }
}