cl-traits-derive 0.0.1

extern crate proc_macro;

use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, quote_spanned};
use syn::{
  parse_macro_input, parse_quote, spanned::Spanned, Data, DeriveInput, Fields, GenericParam,
  Generics, Index,
};

fn add_trait_bounds(trait_name: &Ident, 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!(::cl_traits::#trait_name));
    }
  }
  generics
}

fn fields_fns_with_inputs(trait_fn_name: &Ident, data: &Data) -> TokenStream {
  match *data {
    Data::Struct(ref data) => match data.fields {
      Fields::Named(ref fields) => {
        let fields_fns = fields.named.iter().enumerate().map(|(idx, f)| {
          let name = &f.ident;
          quote_spanned!(f.span()=> self.#name.#trait_fn_name(input.#idx))
        });
        quote!((#(#fields_fns,)*))
      }
      Fields::Unnamed(ref fields) => {
        let fields_fns = fields.unnamed.iter().enumerate().map(|(idx, f)| {
          let idx = Index::from(idx);
          quote_spanned!(f.span()=> self.#idx.#trait_fn_name(input.#idx))
        });
        quote!((#(#fields_fns,)*))
      }
      Fields::Unit => {
        unimplemented!();
      }
    },
    Data::Enum(_) | Data::Union(_) => unimplemented!(),
  }
}

fn fields_fns_without_inputs(trait_fn_name: &Ident, data: &Data) -> TokenStream {
  match *data {
    Data::Struct(ref data) => match data.fields {
      Fields::Named(ref fields) => {
        let fields_fns = fields.named.iter().map(|f| {
          let name = &f.ident;
          quote_spanned!(f.span()=> self.#name.#trait_fn_name())
        });
        quote!((#(#fields_fns,)*))
      }
      Fields::Unnamed(ref fields) => {
        let fields_fns = fields.unnamed.iter().enumerate().map(|(idx, f)| {
          let idx = Index::from(idx);
          quote_spanned!(f.span()=> self.#idx.#trait_fn_name())
        });
        quote!((#(#fields_fns,)*))
      }
      Fields::Unit => {
        unimplemented!();
      }
    },
    Data::Enum(_) | Data::Union(_) => unimplemented!(),
  }
}

fn fields_types(trait_name: &Ident, trait_assoc_ret_name: &Ident, data: &Data) -> TokenStream {
  match *data {
    Data::Struct(ref data) => match data.fields {
      Fields::Named(ref fields) => {
        let field_types = fields.named.iter().map(|f| {
          let ty = &f.ty;
          quote_spanned!(f.span()=> <#ty as #trait_name>::#trait_assoc_ret_name)
        });
        quote!((#(#field_types,)*))
      }
      Fields::Unnamed(ref fields) => {
        let field_types = fields.unnamed.iter().map(|f| {
          let ty = &f.ty;
          quote_spanned!(f.span()=> <#ty as #trait_name>::#trait_assoc_ret_name)
        });
        quote!((#(#field_types,)*))
      }
      Fields::Unit => {
        unimplemented!();
      }
    },
    Data::Enum(_) | Data::Union(_) => unimplemented!(),
  }
}

macro_rules! trait_with_assoc_ret {
  (
    $derive_name:ident,
    $derive_fn_name:ident,
    $trait_name:literal,
    $trait_assoc_ret_name:literal,
    $trait_fn_name:literal,
    $self:ty
  ) => {
    #[proc_macro_derive($derive_name)]
    pub fn $derive_fn_name(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
      let input = parse_macro_input!(input as DeriveInput);

      let data_name = input.ident;
      let trait_name = Ident::new($trait_name, Span::call_site());
      let trait_fn_name = Ident::new($trait_fn_name, Span::call_site());
      let trait_assoc_ret_name = Ident::new($trait_assoc_ret_name, Span::call_site());

      let generics = add_trait_bounds(&trait_name, input.generics);
      let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
      let fields_fns = fields_fns_without_inputs(&trait_fn_name, &input.data);
      let fields_types = fields_types(&trait_name, &trait_assoc_ret_name, &input.data);

      proc_macro::TokenStream::from(quote! {
        impl #impl_generics ::cl_traits::#trait_name for #data_name #ty_generics #where_clause {
          type #trait_assoc_ret_name = #fields_types;

          fn #trait_fn_name(self: $self) -> Self::#trait_assoc_ret_name {
            (#fields_fns)
          }
        }
      })
    }
  };
}

macro_rules! trait_with_assoc_input_and_ret {
  (
    $derive_name:ident,
    $derive_fn_name:ident,
    $trait_name:literal,
    $trait_assoc_input_name:literal,
    $trait_assoc_ret_name:literal,
    $trait_fn_name:literal,
    $self:ty
  ) => {
    #[proc_macro_derive($derive_name)]
    pub fn $derive_fn_name(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
      let input = parse_macro_input!(input as DeriveInput);

      let data_name = input.ident;
      let trait_name = Ident::new($trait_name, Span::call_site());
      let trait_fn_name = Ident::new($trait_fn_name, Span::call_site());
      let trait_assoc_input_name = Ident::new($trait_assoc_input_name, Span::call_site());
      let trait_assoc_ret_name = Ident::new($trait_assoc_ret_name, Span::call_site());

      let generics = add_trait_bounds(&trait_name, input.generics);
      let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
      let fields_fns = fields_fns_with_inputs(&trait_fn_name, &input.data);
      let fields_input_types = fields_types(&trait_name, &trait_assoc_input_name, &input.data);
      let fields_ret_types = fields_types(&trait_name, &trait_assoc_ret_name, &input.data);

      proc_macro::TokenStream::from(quote! {
        impl #impl_generics ::cl_traits::#trait_name for #data_name #ty_generics #where_clause {
          type #trait_assoc_input_name = #fields_input_types;
          type #trait_assoc_ret_name = #fields_ret_types;

          fn #trait_fn_name(self: $self, input: Self::#trait_assoc_input_name) -> Self::#trait_assoc_ret_name {
            #fields_fns
          }
        }
      })
    }
  }
}

trait_with_assoc_ret!(
  WithCapacity,
  derive_with_capacity,
  "Capacity",
  "CapacityRet",
  "capacity",
  &Self
);
trait_with_assoc_ret!(
  WithClear,
  derive_with_clear,
  "Clear",
  "ClearRet",
  "clear",
  &mut Self
);
trait_with_assoc_ret!(
  WithLength,
  derive_with_length,
  "Length",
  "LengthRet",
  "length",
  &Self
);

trait_with_assoc_input_and_ret!(
  WithSwap,
  derive_with_swap,
  "Swap",
  "SwapInput",
  "SwapRet",
  "swap",
  &mut Self
);
trait_with_assoc_input_and_ret!(
  WithTruncate,
  derive_with_truncate,
  "Truncate",
  "TruncateInput",
  "TruncateRet",
  "truncate",
  &mut Self
);