mlua_extras_derive/
lib.rs

1#[macro_use]
2extern crate quote;
3
4use proc_macro::TokenStream;
5use proc_macro2::{Span, TokenStream as TokenStream2};
6use proc_macro_error::{proc_macro_error, abort};
7use syn::spanned::Spanned;
8use venial::{parse_item, Fields, Item};
9
10#[proc_macro_error]
11#[proc_macro_derive(UserData)]
12pub fn derive_user_data(input: TokenStream) -> TokenStream {
13    let input = TokenStream2::from(input);
14    let name = match parse_item(input.clone()) {
15        Ok(Item::Struct(struct_type)) => {
16            struct_type.name.clone()
17        },
18        Ok(Item::Enum(enum_type)) => {
19            enum_type.name.clone()
20        },
21        Err(err) => abort!(err.span(), "{}", err),
22        _ => abort!(input.span(), "only `struct` and `enum` types are supported for TypedUserData")
23    };
24
25    quote!(
26        impl mlua_extras::mlua::UserData for #name {
27            fn add_fields<'lua, F: mlua_extras::mlua::UserDataFields<'lua, Self>>(fields: &mut F) {
28                let mut wrapper = mlua_extras::typed::WrappedGenerator::new(fields);
29                <#name as mlua_extras::typed::TypedUserData>::add_fields(&mut wrapper);
30            }
31
32            fn add_methods<'lua, M: mlua_extras::mlua::UserDataMethods<'lua, Self>>(methods: &mut M) {
33                let mut wrapper = mlua_extras::typed::WrappedGenerator::new(methods);
34                <#name as mlua_extras::typed::TypedUserData>::add_methods(&mut wrapper);
35            }
36        }
37    ).into()
38}
39
40#[proc_macro_error]
41#[proc_macro_derive(Typed, attributes(typed))]
42pub fn derive_typed(input: TokenStream) -> TokenStream {
43    let input = TokenStream2::from(input);
44    match parse_item(input.clone()) {
45        Ok(Item::Struct(struct_type)) => {
46            let name = struct_type.name.clone();
47            let value = syn::LitStr::new(name.to_string().as_str(), Span::call_site());
48            quote!(
49                impl mlua_extras::typed::Typed for #name {
50                    fn ty() -> mlua_extras::typed::Type {
51                        mlua_extras::typed::Type::single(#value)
52                    }
53                }
54            )
55        },
56        Ok(Item::Enum(enum_type)) => {
57            let variants = enum_type.variants
58                .iter()
59                .map(|(variant, _punc)| {
60                    let name = format!("\"{}\"", variant.name);
61                    match &variant.fields {
62                        Fields::Unit => quote!{ mlua_extras::typed::Type::single(#name) },
63                        Fields::Tuple(tf) => {
64                            let tuple_values = tf.fields.iter().map(|(field, _)| {
65                                let ty = field.ty.clone();
66                                quote!{ <#ty as mlua_extras::typed::Typed>::ty() }
67                            }).collect::<Vec<_>>();
68
69                            if tuple_values.len() == 1 {
70                                let first = tuple_values.first().unwrap();
71                                quote!{ #first }
72                            } else {
73                                quote!{ mlua_extras::typed::Type::Tuple(Vec::from([
74                                        #(#tuple_values,)*
75                                ])) }
76                            }
77                        },
78                        Fields::Named(named) => {
79                            let tuple_values = named.fields.iter().map(|(field, _)| {
80                                let name = field.name.to_string();
81                                let ty = field.ty.clone();
82                                quote!{ (#name, <#ty as mlua_extras::typed::Typed>::ty()) }
83                            }).collect::<Vec<_>>();
84                            quote!{ mlua_extras::typed::Type::Struct(std::collections::HashMap::from([
85                                    #(#tuple_values,)*
86                            ])) }
87                        }
88                    }
89                    
90                })
91                .collect::<Vec<_>>();
92
93            // TODO: This should be a union alias
94            let name = enum_type.name.clone();
95            let value = name.to_string();
96            quote!(
97                impl mlua_extras::typed::Typed for #name {
98                    fn ty() -> mlua_extras::typed::Type {
99                        mlua_extras::typed::Type::r#enum(
100                            #value,
101                            [ #(#variants,)* ]
102                        )
103                    }
104                }
105            )
106        },
107        Err(err) => abort!(err.span(), "{}", err),
108        _ => abort!(input.span(), "only `struct` and `enum` types are supported for Typed")
109    }.into()
110}