diff options
Diffstat (limited to 'src')
-rw-r--r-- | src/lib.rs | 97 |
1 files changed, 97 insertions, 0 deletions
diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..8a030c4 --- /dev/null +++ b/src/lib.rs | |||
@@ -0,0 +1,97 @@ | |||
1 | extern crate proc_macro; | ||
2 | |||
3 | use proc_macro::TokenStream; | ||
4 | use quote::{format_ident, quote}; | ||
5 | use syn::punctuated::Punctuated; | ||
6 | use syn::{parse_macro_input, Block, FnArg, ItemFn, Pat, ReturnType, Type}; | ||
7 | |||
8 | #[proc_macro_attribute] | ||
9 | pub fn curry(_attr: TokenStream, item: TokenStream) -> TokenStream { | ||
10 | let parsed = parse_macro_input!(item as ItemFn); | ||
11 | generate_curry(parsed).into() | ||
12 | } | ||
13 | |||
14 | fn extract_type(a: FnArg) -> Box<Type> { | ||
15 | match a { | ||
16 | FnArg::Typed(p) => p.ty, | ||
17 | _ => panic!("Not supported on types with `self`!"), | ||
18 | } | ||
19 | } | ||
20 | |||
21 | fn extract_arg_pat(a: FnArg) -> Box<Pat> { | ||
22 | match a { | ||
23 | FnArg::Typed(p) => p.pat, | ||
24 | _ => panic!("Not supported on types with `self`!"), | ||
25 | } | ||
26 | } | ||
27 | |||
28 | fn extract_return_type(a: ReturnType) -> Box<Type> { | ||
29 | match a { | ||
30 | ReturnType::Type(_, p) => p, | ||
31 | _ => panic!("Not supported on functions without return types!"), | ||
32 | } | ||
33 | } | ||
34 | |||
35 | fn extract_arg_types(fn_args: Punctuated<FnArg, syn::token::Comma>) -> Vec<Box<Type>> { | ||
36 | return fn_args.into_iter().map(extract_type).collect::<Vec<_>>(); | ||
37 | } | ||
38 | |||
39 | fn extract_arg_idents(fn_args: Punctuated<FnArg, syn::token::Comma>) -> Vec<Box<Pat>> { | ||
40 | return fn_args.into_iter().map(extract_arg_pat).collect::<Vec<_>>(); | ||
41 | } | ||
42 | |||
43 | fn generate_type_aliases( | ||
44 | fn_arg_types: &[Box<Type>], | ||
45 | fn_return_type: Box<Type>, | ||
46 | fn_name: &syn::Ident, | ||
47 | ) -> Vec<proc_macro2::TokenStream> { | ||
48 | let type_t0 = format_ident!("_{}_T0", fn_name); | ||
49 | let mut type_aliases = vec![quote! { type #type_t0 = #fn_return_type }]; | ||
50 | for (i, t) in (1..).zip(fn_arg_types.into_iter().rev()) { | ||
51 | let p = format_ident!("_{}_{}", fn_name, format!("T{}", i - 1)); | ||
52 | let n = format_ident!("_{}_{}", fn_name, format!("T{}", i)); | ||
53 | type_aliases.push(quote! { | ||
54 | type #n = impl Fn(#t) -> #p | ||
55 | }) | ||
56 | } | ||
57 | return type_aliases; | ||
58 | } | ||
59 | |||
60 | fn generate_body(fn_args: &[Box<Pat>], body: Box<Block>) -> proc_macro2::TokenStream { | ||
61 | quote! { | ||
62 | return #( move |#fn_args| )* #body | ||
63 | } | ||
64 | } | ||
65 | |||
66 | fn generate_curry(parsed: ItemFn) -> proc_macro2::TokenStream { | ||
67 | let fn_body = parsed.block; | ||
68 | let sig = parsed.sig; | ||
69 | let vis = parsed.vis; | ||
70 | let fn_name = sig.ident; | ||
71 | let fn_args = sig.inputs; | ||
72 | let fn_return_type = sig.output; | ||
73 | |||
74 | let arg_idents = extract_arg_idents(fn_args.clone()); | ||
75 | let first_ident = &arg_idents.first().unwrap(); | ||
76 | let curried_body = generate_body(&arg_idents[1..], fn_body.clone()); | ||
77 | |||
78 | let arg_types = extract_arg_types(fn_args.clone()); | ||
79 | let first_type = &arg_types.first().unwrap(); | ||
80 | let type_aliases = generate_type_aliases( | ||
81 | &arg_types[1..], | ||
82 | extract_return_type(fn_return_type), | ||
83 | &fn_name, | ||
84 | ); | ||
85 | |||
86 | let return_type = format_ident!("_{}_{}", &fn_name, format!("T{}", type_aliases.len() - 1)); | ||
87 | |||
88 | let gen = quote! { | ||
89 | #(#type_aliases);* ; | ||
90 | #vis fn #fn_name (#first_ident: #first_type) -> #return_type { | ||
91 | #curried_body ; | ||
92 | } | ||
93 | }; | ||
94 | |||
95 | println!("{}", gen); | ||
96 | return gen; | ||
97 | } | ||