collum-derive 0.1.0

use darling::{FromDeriveInput, FromMeta};
use proc_macro2::{Span, TokenStream};
use quote::{quote, quote_spanned};
use syn::spanned::Spanned;
use syn::{parse_macro_input, parse_quote, Field, FieldsNamed, FieldsUnnamed, Type};
use syn::{Data, DeriveInput, Fields, GenericParam, Generics, Ident};

#[derive(Copy, Clone, Default, Debug, PartialEq, Eq, FromMeta)]
enum Endian {
    #[default]
    Native,
    Little,
    Big,
}

impl Endian {
    fn transform_captured_type(
        &self,
        ident: Option<&Ident>,
        ty: &Type,
    ) -> (TokenStream, TokenStream) {
        match self {
            Endian::Native => (quote!(#ident), quote!(#ty)),
            Endian::Little => (
                quote!(collum::endian::Little(#ident)),
                quote!(collum::endian::Little<#ty>),
            ),
            Endian::Big => (
                quote!(collum::endian::Big(#ident)),
                quote!(collum::endian::Big<#ty>),
            ),
        }
    }

    fn field_to_captured_type(&self, field: &Field) -> (TokenStream, TokenStream) {
        let ident = &field.ident;
        let ty = &field.ty;

        self.transform_captured_type(ident.as_ref(), ty)
    }
}

#[derive(FromDeriveInput, Default)]
#[darling(default, attributes(collum))]
struct Opts {
    #[darling(default)]
    endian: Endian,
}

fn emit_impl_body_struct_named(
    ident: &Ident,
    accvar: &Ident,
    fields: &FieldsNamed,
    opts: &Opts,
) -> TokenStream {
    let binding: Vec<_> = fields.named.iter().map(|f| &f.ident).collect();

    let (capturing, ty): (Vec<_>, Vec<_>) = fields
        .named
        .iter()
        .map(|f| opts.endian.field_to_captured_type(f))
        .unzip();

    quote! {
        #( let #binding; )*

        #( (#capturing, #accvar) = <#ty as collum::Collum>::take_from_bit_slice(#accvar)?; )*

        Some(( #ident { #( #binding ),* }, #accvar))
    }
}

fn emit_impl_body_struct_unnamed(
    ident: &Ident,
    accvar: &Ident,
    fields: &FieldsUnnamed,
    opts: &Opts,
) -> TokenStream {
    let unnamed_to_numbered = |i: usize, f: &Field| -> Ident {
        let name = format!("v{i}");
        Ident::new(&name, f.span())
    };

    let binding: Vec<_> = fields
        .unnamed
        .iter()
        .enumerate()
        .map(|(i, f)| unnamed_to_numbered(i, f))
        .collect();

    let (capturing, ty): (Vec<_>, Vec<_>) = fields
        .unnamed
        .iter()
        .enumerate()
        .map(|(i, f)| {
            let named = unnamed_to_numbered(i, f);
            opts.endian.transform_captured_type(Some(&named), &f.ty)
        })
        .unzip();

    quote! {
        #( let #binding; )*

        #( (#capturing, #accvar) = <#ty as collum::Collum>::take_from_bit_slice(#accvar)?; )*

        Some((#ident( #( #binding ),* ), #accvar))
    }
}

fn emit_impl_body(ident: &Ident, accvar: &Ident, data: &Data, opts: &Opts) -> TokenStream {
    match *data {
        Data::Struct(ref data) => match &data.fields {
            Fields::Named(fields) => emit_impl_body_struct_named(ident, accvar, fields, opts),
            Fields::Unnamed(fields) => emit_impl_body_struct_unnamed(ident, accvar, fields, opts),
            Fields::Unit => {
                quote!(Some((#ident, #accvar)))
            }
        },
        Data::Enum(..) | Data::Union(..) => unimplemented!(),
    }
}

fn impl_collum(ast: DeriveInput, opts: Opts) -> TokenStream {
    let name = ast.ident;
    let slice_var = Ident::new("slice", Span::call_site());
    let generics = add_collum_sized_bound(ast.generics);

    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    let impl_body = emit_impl_body(&name, &slice_var, &ast.data, &opts);

    let expanded = quote! {
        impl #impl_generics Collum for #name #ty_generics #where_clause {
            fn take_from_bit_slice(mut slice: collum::BitSlice) -> Option<(Self, collum::BitSlice)> {
                #impl_body
            }
        }
    };

    expanded
}

fn impl_collum_sized(ast: DeriveInput) -> TokenStream {
    let name = ast.ident;
    let generics = add_collum_sized_bound(ast.generics);

    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    fn compute_size(data: &Data) -> TokenStream {
        match *data {
            Data::Struct(ref data) => {
                let size = data.fields.iter().map(|f| {
                    let ty = &f.ty;
                    quote_spanned! { f.span() =>
                        <#ty as collum::Sized>::bits()
                    }
                });

                quote! {
                    0 #(+ #size)*
                }
            }
            Data::Enum(..) | Data::Union(..) => unimplemented!(),
        }
    }

    let size = compute_size(&ast.data);

    let expanded = quote! {
        impl #impl_generics collum::Sized for #name #ty_generics #where_clause {
            fn bits() -> usize {
                #size
            }
        }
    };

    expanded
}

fn add_collum_sized_bound(mut generics: Generics) -> Generics {
    for param in &mut generics.params {
        if let GenericParam::Type(ref mut type_param) = *param {
            type_param.bounds.push(parse_quote!(collum::Sized))
        }
    }
    generics
}

#[proc_macro_derive(Collum, attributes(collum))]
pub fn collum_derive(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast = parse_macro_input!(item as DeriveInput);
    let opts = Opts::from_derive_input(&ast).expect("Wrong options");

    proc_macro::TokenStream::from(impl_collum(ast, opts))
}

#[proc_macro_derive(Sized)]
pub fn collum_sized_derive(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let ast = parse_macro_input!(item as DeriveInput);

    proc_macro::TokenStream::from(impl_collum_sized(ast))
}