pmrpc/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{
4    parse::{Parse, ParseStream},
5    parse_macro_input, Ident, Token, Type,
6};
7
8struct Request {
9    request_type: Ident,
10    response_type: Type,
11}
12
13struct DeriveList {
14    derives: Vec<Ident>,
15}
16
17struct RequestList {
18    derives: Option<DeriveList>,
19    requests: Vec<Request>,
20}
21
22impl Parse for DeriveList {
23    fn parse(input: ParseStream) -> syn::Result<Self> {
24        let content;
25        syn::parenthesized!(content in input);
26
27        let mut derives = Vec::new();
28        while !content.is_empty() {
29            derives.push(content.parse()?);
30            if !content.is_empty() {
31                content.parse::<Token![,]>()?;
32            }
33        }
34
35        Ok(DeriveList { derives })
36    }
37}
38
39impl Parse for RequestList {
40    fn parse(input: ParseStream) -> syn::Result<Self> {
41        let derives = if input.peek(syn::token::Paren) {
42            let derives: DeriveList = input.parse()?;
43            Some(derives)
44        } else {
45            None
46        };
47
48        let mut requests = Vec::new();
49
50        while !input.is_empty() {
51            requests.push(input.parse()?);
52            if !input.is_empty() {
53                input.parse::<Token![,]>()?;
54            }
55        }
56
57        Ok(RequestList { derives, requests })
58    }
59}
60
61impl Parse for Request {
62    fn parse(input: ParseStream) -> syn::Result<Self> {
63        let request_type = input.parse()?;
64        input.parse::<Token![=>]>()?;
65
66        // Parse the entire response type as a single Type
67        let response_type = input.parse()?;
68
69        Ok(Request {
70            request_type,
71            response_type,
72        })
73    }
74}
75
76#[proc_macro]
77pub fn define_requests(input: TokenStream) -> TokenStream {
78    let RequestList { derives, requests } = parse_macro_input!(input as RequestList);
79
80    // Create the derive attribute if specified
81    let derive_attr = derives.map(|d| {
82        let derives = d.derives;
83        quote! { #[derive(#(#derives),*)] }
84    });
85
86    let responds_with_impls = requests.iter().map(|req| {
87        let name = &req.request_type;
88        let response_type = &req.response_type;
89
90        quote! {
91            impl RespondsWith<#response_type> for #name {
92                fn to_enum(self) -> Requests {
93                    Requests::#name(self)
94                }
95
96                fn resp_enum(r: #response_type) -> Responses {
97                    Responses::#response_type(r)
98                }
99
100                fn resp_from_enum(r: Responses) -> #response_type {
101                    match r {
102                        Responses::#response_type(x) => x,
103                        _ => panic!("broken code"),
104                    }
105                }
106
107                fn name() -> &'static str {
108                    stringify!(#name)
109                }
110            }
111        }
112    });
113
114    let request_enum_variants = requests.iter().map(|req| {
115        let name = &req.request_type;
116        quote! {
117            #name(#name)
118        }
119    });
120
121    let response_types: Vec<_> =
122        requests
123            .iter()
124            .map(|req| &req.response_type)
125            .fold(Vec::new(), |mut acc, ty| {
126                if !acc
127                    .iter()
128                    .any(|x: &&Type| quote!(#x).to_string() == quote!(#ty).to_string())
129                {
130                    acc.push(ty);
131                }
132                acc
133            });
134
135    let expanded = quote! {
136        #(#responds_with_impls)*
137
138        pub trait RespondsWith<Resp> {
139            fn to_enum(self) -> Requests;
140            fn resp_enum(r: Resp) -> Responses;
141            fn resp_from_enum(r: Responses) -> Resp;
142            fn name() -> &'static str;
143        }
144
145        #derive_attr
146        #[derive(Debug)]
147        pub enum Requests {
148            #(#request_enum_variants,)*
149        }
150
151        #derive_attr
152        #[derive(Debug)]
153        pub enum Responses {
154            #(#response_types(#response_types),)*
155        }
156    };
157
158    expanded.into()
159}