polybox_codegen/
lib.rs

1//! Macros for the `polybox` crate.
2//!
3//! See [GitHub](https://github.com/jvdwrf/polybox) for more information.
4
5extern crate proc_macro;
6use proc_macro::TokenStream;
7use quote::quote;
8use syn::{parse_macro_input, Data, DeriveInput, Expr, Fields, Lit, Type};
9
10/// Derives the `Interface` trait for an enum, allowing it to be used as a message
11/// interface in the Polybox framework.
12///
13/// Under the hood, this macro generates implementations for the `Interface`, `Message`, and `AsSet` traits,
14/// as well as `FromPayload` and `TryIntoPayload` for each variant of the enum.
15///
16/// The macro expects the enum variants to be of the form `Variant(Payload<T>)`,
17/// where `T` is a type that implements the `Message` trait.
18#[proc_macro_derive(Interface, attributes(polybox))]
19pub fn derive_interface(input: TokenStream) -> TokenStream {
20    let input = parse_macro_input!(input as DeriveInput);
21    let enum_name = &input.ident;
22
23    // 1. Determine the base path (default: ::polybox)
24    let mut base_path: syn::Path = syn::parse_str("::polybox").unwrap();
25
26    for attr in &input.attrs {
27        if attr.path().is_ident("polybox") {
28            // Correct way to parse nested meta (e.g. #[polybox(crate = "...")] ) in syn 2.0
29            let _ = attr.parse_nested_meta(|meta| {
30                if meta.path.is_ident("crate") {
31                    let value = meta.value()?;
32                    let expr: Expr = value.parse()?;
33                    if let Expr::Lit(syn::ExprLit {
34                        lit: Lit::Str(lit_str),
35                        ..
36                    }) = expr
37                    {
38                        if let Ok(parsed_path) = syn::parse_str::<syn::Path>(&lit_str.value()) {
39                            base_path = parsed_path;
40                        }
41                    }
42                }
43                Ok(())
44            });
45        }
46    }
47
48    // Ensure we are working with an enum
49    let variants = match &input.data {
50        Data::Enum(data_enum) => &data_enum.variants,
51        _ => panic!("Interface derive can only be used on enums"),
52    };
53
54    let mut inner_types = Vec::new();
55    let mut try_from_matches = Vec::new();
56    let mut try_into_matches = Vec::new();
57    let mut into_matches = Vec::new();
58    let mut from_impls = Vec::new();
59
60    for variant in variants {
61        let variant_name = &variant.ident;
62
63        match &variant.fields {
64            Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
65                let field_type = &fields.unnamed[0].ty;
66
67                let inner_type = extract_inner_type(field_type)
68                    .expect("Interface variants must be of type Payload<T>");
69
70                inner_types.push(inner_type);
71
72                try_from_matches.push(quote! {
73                    let payload = match payload.downcast::<#inner_type>() {
74                        Ok(payload) => return Ok(Self::#variant_name(payload)),
75                        Err(payload) => payload,
76                    };
77                });
78
79                try_into_matches.push(quote! {
80                    if id == std::any::TypeId::of::<#inner_type>() {
81                        if let Self::#variant_name(payload) = self {
82                            // SAFETY: Verified type matches dynamic I parameter.
83                            let converted = unsafe {
84                                std::mem::transmute_copy::<#base_path::Payload<#inner_type>, #base_path::Payload<I>>(&payload)
85                            };
86                            std::mem::forget(payload);
87                            return Ok(converted);
88                        }
89                    }
90                });
91
92                into_matches.push(quote! {
93                    Self::#variant_name(payload) => #base_path::BoxedPayload::new::<#inner_type>(payload),
94                });
95
96                from_impls.push(quote! {
97                    impl #base_path::FromPayload<#inner_type> for #enum_name {
98                        fn from_payload(payload: #base_path::Payload<#inner_type>) -> Self {
99                            Self::#variant_name(payload)
100                        }
101                    }
102
103                    impl #base_path::TryIntoPayload<#inner_type> for #enum_name {
104                        fn try_into_payload(self) -> Result<#base_path::Payload<#inner_type>, Self> {
105                            if let #enum_name::#variant_name(payload) = self {
106                                Ok(payload)
107                            } else {
108                                Err(self)
109                            }
110                        }
111                    }
112                });
113            }
114            _ => panic!("Interface derive only supports variants with a single unnamed field, e.g., A(Payload<T>)"),
115        }
116    }
117
118    let expanded = quote! {
119        impl #base_path::Interface for #enum_name {
120            fn try_from_boxed_payload(payload: #base_path::BoxedPayload) -> Result<Self, #base_path::BoxedPayload> {
121                #(#try_from_matches)*
122                Err(payload)
123            }
124
125            // Could be added to improve performance, but would require unsafe transmute to avoid double downcasting.
126            // fn try_into_payload<I: #base_path::Message>(self) -> Result<#base_path::Payload<I>, Self> {
127            //     let id = std::any::TypeId::of::<I>();
128            //     #(#try_into_matches)*
129            //     Err(self)
130            // }
131
132            fn into_boxed_payload(self) -> #base_path::BoxedPayload {
133                match self {
134                    #(#into_matches)*
135                }
136            }
137        }
138
139        impl #base_path::Message for #enum_name {
140            type Kind = #base_path::FireAndForget;
141        }
142
143        impl #base_path::type_sets::AsSet for #enum_name {
144            type Set = #base_path::type_sets::Set![#(#inner_types),*];
145        }
146
147        impl #base_path::TryIntoPayload<#enum_name> for #enum_name {
148            fn try_into_payload(self) -> Result<#base_path::Payload<#enum_name>, Self> {
149                Ok(self)
150            }
151        }
152
153        impl #base_path::FromPayload<#enum_name> for #enum_name {
154            fn from_payload(payload: #base_path::Payload<#enum_name>) -> Self {
155                payload
156            }
157        }
158
159        #(#from_impls)*
160    };
161
162    TokenStream::from(expanded)
163}
164
165fn extract_inner_type(ty: &Type) -> Option<&Type> {
166    if let Type::Path(type_path) = ty {
167        let segment = type_path.path.segments.last()?;
168        if segment.ident == "Payload" {
169            if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
170                if let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first() {
171                    return Some(inner_ty);
172                }
173            }
174        }
175    }
176    None
177}
178
179/// Derives the `Message` trait for a struct, allowing it to be used as a message
180/// in the Polybox framework.
181///
182/// This macro accepts an optional `reply` attribute to specify the reply type for the message.
183///
184/// # Example
185/// ```ignore
186/// #[derive(Message)]
187/// struct SimpleMessage;
188///
189/// #[derive(Message)]
190/// #[msg(reply = u32)]
191/// struct MessageWithReply;
192/// ```
193#[proc_macro_derive(Message, attributes(polybox, msg))]
194pub fn derive_invocation(input: TokenStream) -> TokenStream {
195    let input = parse_macro_input!(input as DeriveInput);
196    let name = &input.ident;
197
198    // // 1. Determine the base path (default: ::polybox)
199    // let mut base_path: syn::Path = syn::parse_str("::polybox").unwrap();
200
201    // // 2. Determine the default Kind (default: ::polybox::FireAndForget)
202    // let mut kind_type = quote!(::polybox::FireAndForget);
203
204    let polybox_attr = &input
205        .attrs
206        .iter()
207        .find(|attr| attr.path().is_ident("polybox"));
208
209    let base_path = if let Some(attr) = polybox_attr {
210        let mut base_path: syn::Path = syn::parse_str("::polybox").unwrap();
211        let _ = attr.parse_nested_meta(|meta| {
212            if meta.path.is_ident("crate") {
213                let value = meta.value()?;
214                let expr: Expr = value.parse()?;
215                if let Expr::Lit(syn::ExprLit {
216                    lit: Lit::Str(lit_str),
217                    ..
218                }) = expr
219                {
220                    if let Ok(parsed_path) = syn::parse_str::<syn::Path>(&lit_str.value()) {
221                        base_path = parsed_path;
222                    }
223                }
224            }
225            Ok(())
226        });
227        base_path
228    } else {
229        syn::parse_str("::polybox").unwrap()
230    };
231
232    let invoke_attr = &input.attrs.iter().find(|attr| attr.path().is_ident("msg"));
233
234    let kind_type = if let Some(attr) = invoke_attr {
235        let mut kind_type = quote!(#base_path::FireAndForget);
236        let _ = attr.parse_nested_meta(|meta| {
237            if meta.path.is_ident("reply") {
238                let value = meta.value()?;
239                if let Ok(parsed_type) = value.parse::<Type>() {
240                    kind_type = quote! {
241                        #base_path::Request<#parsed_type>
242                    }
243                }
244            } else {
245                panic!("Only `reply` is expected")
246            }
247            Ok(())
248        });
249        kind_type
250    } else {
251        quote!(#base_path::FireAndForget)
252    };
253
254    // Handle generics if the struct/enum is generic
255    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
256
257    let expanded = quote! {
258        impl #impl_generics #base_path::Message for #name #ty_generics #where_clause {
259            type Kind = #kind_type;
260        }
261    };
262
263    TokenStream::from(expanded)
264}
polybox_codegen/lib.rs

polybox_codegen/
lib.rs
