enum-repr-conv 0.1.0

use proc_macro2::TokenStream;
use quote::{quote, ToTokens};
use syn::{parse_quote, Ident};

use crate::transform::{Transformed, ValueExpr, ValueExprMatch};

pub(crate) fn produce(transformed: Transformed) -> TokenStream {
    let Transformed {
        repr_type,
        variants,
        mut item,
    } = transformed;

    item.attrs.push(parse_quote!(#[repr(#repr_type)]));

    struct TryFromPattern<'a> {
        repr_type: &'a Ident,
        value_expr: ValueExpr,
    }

    impl<'a> ToTokens for TryFromPattern<'a> {
        fn to_tokens(&self, tokens: &mut TokenStream) {
            let TryFromPattern {
                repr_type,
                value_expr,
            } = self;
            let ValueExpr { base, offset } = value_expr;
            let pat = match (base, offset) {
                (None, offset) => quote!(val if val == #offset as #repr_type),
                (Some(expr), 0) => quote!(val if val == #expr),
                (Some(expr), offset) => {
                    quote!(val if val == #repr_type::wrapping_add(#expr, #offset as #repr_type))
                }
            };
            pat.to_tokens(tokens);
        }
    }

    struct TryFromMatchArm<'a, 'b> {
        pattern: TryFromPattern<'a>,
        enum_ident: &'b Ident,
        variant: Ident,
    }

    impl<'a, 'b> ToTokens for TryFromMatchArm<'a, 'b> {
        fn to_tokens(&self, tokens: &mut TokenStream) {
            let TryFromMatchArm {
                pattern,
                enum_ident,
                variant,
            } = self;

            let arm = quote!(#pattern => Ok(#enum_ident :: #variant));
            arm.to_tokens(tokens);
        }
    }

    let enum_ident = &item.ident;

    let conversions = variants
        .into_iter()
        .map(|ValueExprMatch(ident, value_expr)| {
            let pattern = TryFromPattern {
                repr_type: &repr_type,
                value_expr,
            };
            let arm = TryFromMatchArm {
                pattern,
                enum_ident,
                variant: ident,
            };
            quote!(#arm)
        })
        .collect::<Vec<_>>();

    quote!(
        #item

        impl ::std::convert::From<#enum_ident> for #repr_type {
            fn from(e: #enum_ident) -> Self {
                e as Self
            }
        }

        impl ::std::convert::From<&#enum_ident> for #repr_type {
            fn from(e: &#enum_ident) -> Self {
                unsafe { *::std::mem::transmute::<&#enum_ident, &#repr_type>(e) }
            }
        }

        impl ::std::convert::TryFrom<#repr_type> for #enum_ident {
            type Error = ();

            fn try_from(value: #repr_type) -> Result<Self, Self::Error> {
                match value {
                    #(#conversions,)*
                    _ => Err(())
                }
            }
        }

        impl ::std::convert::TryFrom<&#repr_type> for #enum_ident {
            type Error = ();

            fn try_from(value: &#repr_type) -> Result<Self, Self::Error> {
                #enum_ident::try_from(*value)
            }
        }
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use syn::{parse::Parse, parse_quote, ItemEnum, ItemImpl};
    struct EnumDefAndImplBlocks {
        #[allow(dead_code)]
        enum_def: ItemEnum,
        #[allow(dead_code)]
        impl_blocks: [ItemImpl; 4],
    }

    impl Parse for EnumDefAndImplBlocks {
        fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
            let enum_def: ItemEnum = input.parse()?;
            Ok(Self {
                enum_def,
                impl_blocks: [
                    input.parse()?,
                    input.parse()?,
                    input.parse()?,
                    input.parse()?,
                ],
            })
        }
    }

    #[test]
    fn output_is_enum_def_and_impl_blocks() {
        let repr_type: Ident = parse_quote!(u16);

        let variants = vec![
            ValueExprMatch(
                parse_quote!(Zeroth),
                ValueExpr {
                    base: None,
                    offset: 0,
                },
            ),
            ValueExprMatch(
                parse_quote!(First),
                ValueExpr {
                    base: Some(parse_quote! {2 * 7}),
                    offset: 0,
                },
            ),
            ValueExprMatch(
                parse_quote!(Second),
                ValueExpr {
                    base: Some(parse_quote! {2 * 7}),
                    offset: 1,
                },
            ),
            ValueExprMatch(
                parse_quote!(Third),
                ValueExpr {
                    base: Some(parse_quote! {1}),
                    offset: 0,
                },
            ),
        ];

        let item: ItemEnum = parse_quote!(
            enum Fieldless {
                Zeroth,
                First = 2 * 7,
                Second,
                Third = 1,
            }
        );

        let transformed = Transformed {
            repr_type,
            variants,
            item,
        };

        let code = produce(transformed);

        let _ = syn::parse2::<EnumDefAndImplBlocks>(code).unwrap();
    }
}