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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157

#![feature(iter_array_chunks)]
#![feature(let_chains)]

use convert_case::{Case, Casing};
use proc_macro::TokenStream as TokenStreamV1;
use proc_macro2::{Delimiter, Ident, TokenStream as TokenStreamV2, TokenTree};
use quote::{format_ident, quote};

fn get_fn_name(mut tokens_iter: impl Iterator<Item = TokenTree>) -> Option<Ident> {
    if let TokenTree::Ident(ident) = tokens_iter.nth(1)? {
        Some(ident)
    } else {
        None
    }
}

fn get_fn_return_type(tokens_iter: impl Iterator<Item = TokenTree>) -> Option<Ident> {
    let mut is_next_token_return_type = false;
    let return_type = tokens_iter.array_chunks::<2>().find(|[token1, token2]| {
        if is_next_token_return_type {
            return true;
        }
        if let TokenTree::Punct(punct1) = token1 && let TokenTree::Punct(punct2) = token2 {
            if punct1.as_char() == '-' && punct2.as_char() == '>' {
                is_next_token_return_type = true;
                return false;
            } else {
                return false;
            }
        }
        false
    })?;

    if let TokenTree::Ident(return_type) = return_type[0].clone() {
        Some(return_type)
    } else {
        None
    }
}

fn generate_struct(
    fn_name: &Ident,
    mut tokens_iter: impl Iterator<Item = TokenTree>,
) -> Option<(TokenStreamV2, Ident)> {
    let struct_name = format_ident!("{}Args", fn_name.to_string().to_case(Case::Pascal));

    let fn_args_tokens = {
        let fn_args_group = tokens_iter.find(|token_tree| {
            if let TokenTree::Group(group) = token_tree {
                group.delimiter() == Delimiter::Parenthesis
            } else {
                false
            }
        })?;
        if let TokenTree::Group(group) = fn_args_group {
            group.stream()
        } else {
            return None;
        }
    };

    Some((
        quote! {
            #[derive(Deserialize, Debug)]
            struct #struct_name {
                #fn_args_tokens
            }
        },
        struct_name,
    ))
}

fn generate_thunk(
    fn_name: &Ident,
    struct_name: &Ident,
    return_type: Option<Ident>,
    mut tokens_iter: impl Iterator<Item = TokenTree>,
) -> Option<TokenStreamV2> {
    let thunk_name = format_ident!("{}_thunk", fn_name);

    let struct_tokens = {
        let struct_group = tokens_iter.find(|token_tree| {
            if let TokenTree::Group(group) = token_tree {
                group.delimiter() == Delimiter::Brace
            } else {
                false
            }
        })?;
        if let TokenTree::Group(group) = struct_group {
            group.stream()
        } else {
            return None;
        }
    };

    let mut should_filter_next = false;
    let unwrap_stream = struct_tokens
        .into_iter()
        .filter(|token_tree| {
            if should_filter_next {
                should_filter_next = false;
                return false;
            }
            if let TokenTree::Punct(punct) = token_tree {
                if punct.as_char() == ':' {
                    should_filter_next = true;
                    return false;
                } else {
                    return true;
                }
            }
            true
        })
        .collect::<TokenStreamV2>();

    if let Some(return_type) = return_type {
        Some(quote! {
            fn #thunk_name (args: #struct_name) -> #return_type {
                let #struct_name { #unwrap_stream } = args;
                #fn_name(#unwrap_stream)
            }
        })
    } else {
        Some(quote! {
            fn #thunk_name (args: #struct_name) {
                let #struct_name { #unwrap_stream } = args;
                #fn_name(#unwrap_stream);
            }
        })
    }
}

#[proc_macro_attribute]
pub fn server_function(_attr: TokenStreamV1, item: TokenStreamV1) -> TokenStreamV1 {
    let item = Into::<TokenStreamV2>::into(item);
    let mut item_iter = item.clone().into_iter();

    let fn_name = get_fn_name(&mut item_iter).expect("Failed to get function name!");
    let (args_struct, args_struct_name) = generate_struct(&fn_name, &mut item_iter)
        .expect("Failed to generate function arguments struct!");
    let return_type = get_fn_return_type(&mut item_iter);
    let thunk = generate_thunk(
        &fn_name,
        &args_struct_name,
        return_type,
        args_struct.clone().into_iter(),
    )
    .expect("Failed to generate function thunk!");

    quote! {
        #args_struct
        #thunk

        #item
    }
    .into()
}