entropic-derive 0.1.1

Traits for converting Rust data structures to/from unstructured bytes
Documentation
use proc_macro2::TokenStream;
use std::iter;
use syn::spanned::Spanned;

pub fn expand_derive_entropic(ast: syn::DeriveInput) -> syn::Result<TokenStream> {
    let mut tokens = check_fields_entropic(&ast)?;

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

    let (from_entropy_tokens, to_entropy_tokens) = match ast.data {
        syn::Data::Enum(enum_ast) => derive_entropic_enum(&ident, enum_ast)?,
        syn::Data::Struct(struct_ast) => derive_entropic_struct(struct_ast)?,
        syn::Data::Union(union_ast) => {
            return Err(syn::Error::new(
                union_ast.union_token.span,
                "Entropic derive unimplemented for Union types",
            ))
        }
    };

    tokens.extend(quote::quote!{
        #[automatically_derived]
        impl #impl_generics entropic::Entropic for #ident #ty_generics #where_clause {
            #[inline]
            fn from_entropy_source<'a, I: Iterator<Item=&'a u8>, E: entropic::scheme::EntropyScheme>(__entropic_source: &mut entropic::Source<'a, I, E>) -> Result<Self, entropic::EntropicError> {
                #from_entropy_tokens
            }

            #[inline]
            fn to_entropy_sink<'a, I: Iterator<Item=&'a mut u8>, E: entropic::scheme::EntropyScheme>(&self, __entropic_sink: &mut entropic::Sink<'a, I, E>) -> Result<usize, entropic::EntropicError> {
                #to_entropy_tokens
            }
        }
    });

    // println!("{}", tokens);

    Ok(tokens)
}

fn check_fields_entropic(ast: &syn::DeriveInput) -> syn::Result<TokenStream> {
    let mut checks_tokens = TokenStream::new();

    match &ast.data {
        syn::Data::Struct(struct_ast) => {
            for (idx, field) in struct_ast.fields.iter().enumerate() {
                let ty = &field.ty;
                let ty_span = extract_type_span(ty)?;

                let assertion_ident =
                    quote::format_ident!("_AssertStruct{}Field{}IsEntropic", ast.ident, idx);

                checks_tokens.extend(quote::quote_spanned! {ty_span=>
                    struct #assertion_ident where #ty: entropic::Entropic;
                });
            }
        }
        syn::Data::Enum(enum_ast) => {
            for variant in &enum_ast.variants {
                for (idx, field) in variant.fields.iter().enumerate() {
                    let ty = &field.ty;
                    let ty_span = extract_type_span(ty)?;

                    let assertion_ident = quote::format_ident!(
                        "_AssertEnum{}Variant{}Field{}IsEntropic",
                        &ast.ident,
                        &variant.ident,
                        idx
                    );

                    checks_tokens.extend(quote::quote_spanned! {ty_span=>
                        struct #assertion_ident where #ty: entropic::Entropic;
                    });
                }
            }
        }
        syn::Data::Union(_) => (),
    }

    Ok(checks_tokens)
}

fn derive_entropic_struct(struct_ast: syn::DataStruct) -> syn::Result<(TokenStream, TokenStream)> {
    match &struct_ast.fields {
        syn::Fields::Named(fields) => {
            let from_entropy_tokens =
                derive_from_entropy_named_fields(quote::quote! { Self }, fields)?;

            let field_idents = extract_named_fields(fields)?;
            let to_entropy_tokens = derive_to_entropy_fields(
                field_idents.into_iter().map(|f| quote::quote! { self.#f }),
            )?;

            Ok((
                from_entropy_tokens,
                quote::quote! { let mut __entropic_length = 0usize; #to_entropy_tokens },
            ))
        }
        syn::Fields::Unnamed(fields) => {
            let from_entropy_tokens =
                derive_from_entropy_unnamed_fields(&quote::quote! { Self }, fields)?;

            let field_idents = (0..fields.unnamed.len()).map(|idx| {
                let struct_index = syn::Index::from(idx);
                quote::quote! { self.#struct_index }
            });
            let to_entropy_tokens = derive_to_entropy_fields(field_idents)?;

            Ok((
                from_entropy_tokens,
                quote::quote! { let mut __entropic_length = 0usize; #to_entropy_tokens },
            ))
        }
        syn::Fields::Unit => {
            let from_entropy_tokens = quote::quote! { Ok(Self) };
            let to_entropy_tokens = quote::quote! { Ok(0) };

            Ok((from_entropy_tokens, to_entropy_tokens))
        }
    }
}

fn derive_from_entropy_named_fields(
    base_ident: TokenStream,
    fields: &syn::FieldsNamed,
) -> syn::Result<TokenStream> {
    let mut field_insts = Vec::new();
    for field in fields.named.iter() {
        let field_ident = field
            .ident
            .as_ref()
            .ok_or(syn::Error::new_spanned(field, "missing identifier name"))?;
        field_insts.push(quote::quote! {
            #field_ident: __entropic_source.get_entropic()?
        });
    }

    Ok(quote::quote! {
        Ok(#base_ident {
            #(#field_insts),*
        })
    })
}

fn derive_to_entropy_fields<I: Iterator<Item = TokenStream>>(
    field_idents: I,
) -> syn::Result<TokenStream> {
    let mut to_entropy_steps = Vec::new();
    for field_ident in field_idents {
        to_entropy_steps.push(quote::quote! {
            __entropic_length = __entropic_length.checked_add(#field_ident.to_entropy_sink(__entropic_sink)?).ok_or(entropic::EntropicError::Internal)?;
        });
    }

    Ok(quote::quote! {
        #(#to_entropy_steps)*
        Ok(__entropic_length)
    })
}

fn derive_from_entropy_unnamed_fields(
    base_ident: &TokenStream,
    fields: &syn::FieldsUnnamed,
) -> syn::Result<TokenStream> {
    let field_inst = quote::quote! {
        __entropic_source.get_entropic()?
    };

    let field_insts = iter::repeat(field_inst).take(fields.unnamed.len());

    Ok(quote::quote! {
        Ok(#base_ident (#(#field_insts),*))
    })
}

fn derive_entropic_enum(
    enum_ident: &syn::Ident,
    enum_ast: syn::DataEnum,
) -> syn::Result<(TokenStream, TokenStream)> {
    let variants_end = match enum_ast.variants.len().checked_sub(1) {
        Some(len) => len,
        None => {
            return Err(syn::Error::new_spanned(
                &enum_ast.variants,
                "Entropic derive not supported for enum with no variants",
            ))
        }
    };

    let mut from_entropy_matches = Vec::new();
    let mut to_entropy_matches = Vec::new();

    for (idx, variant) in enum_ast.variants.iter().enumerate() {
        let ident = &variant.ident;
        let base_variant_ident = quote::quote! { #enum_ident::#ident };
        match &variant.fields {
            syn::Fields::Named(fields) => {
                let field_idents = extract_named_fields(fields)?;

                let from_entropy_tokens =
                    derive_from_entropy_named_fields(quote::quote! { Self::#ident }, fields)?;
                let to_entropy_tokens = derive_to_entropy_fields(
                    field_idents
                        .clone()
                        .into_iter()
                        .map(|i| quote::quote! { #i }),
                )?;

                from_entropy_matches.push(quote::quote! {
                    #idx => {
                        #from_entropy_tokens
                    }
                });

                to_entropy_matches.push(quote::quote!{
                    #base_variant_ident { #(#field_idents),* } => {
                        let mut __entropic_length = __entropic_sink.put_uniform_range(0..=#variants_end, #idx)?;
                        #to_entropy_tokens
                    }
                });
            }
            syn::Fields::Unnamed(fields) => {
                let field_idents: Vec<syn::Ident> = (0..fields.unnamed.len())
                    .map(|idx| quote::format_ident!("__entropic_field_{}", idx))
                    .collect();

                let from_entropy_tokens =
                    derive_from_entropy_unnamed_fields(&base_variant_ident, fields)?;
                let to_entropy_tokens = derive_to_entropy_fields(
                    field_idents
                        .clone()
                        .into_iter()
                        .map(|f| quote::quote! { #f }),
                )?;

                from_entropy_matches.push(quote::quote! {
                    #idx => {
                        #from_entropy_tokens
                    }
                });

                to_entropy_matches.push(quote::quote!{
                    #base_variant_ident ( #(#field_idents),* ) => {
                        let mut __entropic_length = __entropic_sink.put_uniform_range(0..=#variants_end, #idx)?;
                        #to_entropy_tokens
                    }
                });
            }
            syn::Fields::Unit => {
                from_entropy_matches.push(quote::quote! {
                    #idx => Ok(#base_variant_ident),
                });

                to_entropy_matches.push(quote::quote!{
                    #base_variant_ident => Ok(__entropic_sink.put_uniform_range(0..=#variants_end, #idx)?),
                });
            }
        }
    }

    let from_entropy_tokens = quote::quote! {
        let __entropic_variant_idx = __entropic_source.get_uniform_range(0..=#variants_end)?;
        match __entropic_variant_idx {
            #(#from_entropy_matches)*
            _ => Err(entropic::EntropicError::Internal)
        }
    };

    let to_entropy_tokens = quote::quote! {
        match self {
            #(#to_entropy_matches)*
            _ => Err(entropic::EntropicError::Internal)
        }
    };

    Ok((from_entropy_tokens, to_entropy_tokens))
}

fn extract_named_fields(fields: &syn::FieldsNamed) -> syn::Result<Vec<syn::Ident>> {
    let mut field_idents = Vec::new();
    for field in fields.named.iter() {
        field_idents.push(
            field
                .ident
                .as_ref()
                .ok_or(syn::Error::new_spanned(field, "missing identifier name"))?
                .clone(),
        );
    }

    Ok(field_idents)
}

fn extract_type_span(ty: &syn::Type) -> syn::Result<proc_macro2::Span> {
    Ok(match ty {
        syn::Type::Array(t) => t.span(),
        syn::Type::BareFn(t) => t.span(),
        syn::Type::Group(t) => t.span(),
        syn::Type::ImplTrait(t) => t.span(),
        syn::Type::Infer(t) => t.span(),
        syn::Type::Macro(t) => t.span(),
        syn::Type::Never(t) => t.span(),
        syn::Type::Paren(t) => t.span(),
        syn::Type::Path(path) => path
            .path
            .segments
            .last()
            .ok_or(syn::Error::new_spanned(
                path,
                "Entropic derive unimplemented for type path with zero segments",
            ))?
            .span(),
        syn::Type::Ptr(t) => t.span(),
        syn::Type::Reference(t) => t.span(),
        syn::Type::Slice(t) => t.span(),
        syn::Type::TraitObject(t) => t.span(),
        syn::Type::Tuple(t) => t.span(),
        syn::Type::Verbatim(t) => t.span(),
        _ => return Err(syn::Error::new_spanned(ty, "unrecognized type")),
    })
}