1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

///# Masala
///
/// An autocurrying macro for Rust.
///
/// ## Usage
///
/// This crate requires nightly:
///
///```rust
///#![feature(type_alias_impl_trait, min_type_alias_impl_trait)]
///use masala::curry;
///
///#[curry]
///fn add(a: isize, b: isize) -> isize {
///    a + b
///}
///
///fn main() {
///    println!("{}", add(33)(42));
///}
///```
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::punctuated::Punctuated;
use syn::{parse_macro_input, Block, FnArg, ItemFn, Pat, ReturnType, Type};

#[proc_macro_attribute]
pub fn curry(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let parsed = parse_macro_input!(item as ItemFn);
    cook_curry(parsed).into()
}

fn cook_curry(parsed: ItemFn) -> proc_macro2::TokenStream {
    let fn_body = parsed.block;
    let sig = parsed.sig;
    let vis = parsed.vis;
    let fn_name = sig.ident;
    let fn_args = sig.inputs;
    let fn_return_type = sig.output;

    let arg_idents = extract_arg_idents(fn_args.clone());
    let first_ident = &arg_idents.first().unwrap();

    let curried_body = generate_body(&arg_idents[1..], fn_body.clone());

    let arg_types = extract_arg_types(fn_args.clone());
    let first_type = &arg_types.first().unwrap();
    let type_aliases = generate_type_aliases(
        &arg_types[1..],
        extract_return_type(fn_return_type),
        &fn_name,
    );
    let return_type = format_ident!(
        "{}{}",
        title_case(&fn_name.to_string()),
        format!("T{}", type_aliases.len() - 1)
    );

    quote! {
        #(#type_aliases);* ;
        #vis fn #fn_name (#first_ident: #first_type) -> #return_type {
            #curried_body ;
        }
    }
}

fn extract_arg_idents(fn_args: Punctuated<FnArg, syn::token::Comma>) -> Vec<Box<Pat>> {
    fn_args.into_iter().map(extract_arg_pat).collect::<Vec<_>>()
}

fn extract_arg_pat(a: FnArg) -> Box<Pat> {
    match a {
        FnArg::Typed(p) => p.pat,
        _ => panic!("Not supported on types with `self!`"),
    }
}

fn extract_arg_types(fn_args: Punctuated<FnArg, syn::token::Comma>) -> Vec<Box<Type>> {
    fn_args.into_iter().map(extract_type).collect::<Vec<_>>()
}

fn extract_return_type(a: ReturnType) -> Box<Type> {
    match a {
        ReturnType::Type(_, p) => p,
        _ => panic!("Not supported on functions without return types!"),
    }
}

fn extract_type(a: FnArg) -> Box<Type> {
    match a {
        FnArg::Typed(p) => p.ty,
        _ => panic!("Not supported on types with `self!`"),
    }
}

fn generate_body(fn_args: &[Box<Pat>], body: Box<Block>) -> proc_macro2::TokenStream {
    quote! {
        return #( move |#fn_args| )* #body
    }
}

fn generate_type_aliases(
    fn_arg_types: &[Box<Type>],
    fn_return_type: Box<Type>,
    fn_name: &syn::Ident,
) -> Vec<proc_macro2::TokenStream> {
    let type_t0 = format_ident!("{}T0", title_case(&fn_name.to_string()));
    let mut type_aliases = vec![quote! { type #type_t0 = #fn_return_type}];
    for (i, t) in (1..).zip(fn_arg_types.into_iter().rev()) {
        let func = title_case(&fn_name.to_string());
        let p = format_ident!("{}{}", func, format!("T{}", i - 1));
        let n = format_ident!("{}{}", func, format!("T{}", i));
        type_aliases.push(quote! {
            type #n = impl Fn(#t) -> #p
        });
    }
    type_aliases
}

fn title_case(s: &str) -> String {
    let mut c = s.chars();
    match c.next() {
        None => String::new(),
        Some(f) => f.to_uppercase().collect::<String>() + c.as_str(),
    }
}