mlua_derive/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::{Ident, Span};
3use quote::quote;
4use syn::meta::ParseNestedMeta;
5use syn::{parse_macro_input, ItemFn, LitStr, Result};
6
7#[cfg(feature = "macros")]
8use {
9    crate::chunk::Chunk, proc_macro::TokenTree, proc_macro2::TokenStream as TokenStream2,
10    proc_macro_error2::proc_macro_error,
11};
12
13#[derive(Default)]
14struct ModuleAttributes {
15    name: Option<Ident>,
16    skip_memory_check: bool,
17}
18
19impl ModuleAttributes {
20    fn parse(&mut self, meta: ParseNestedMeta) -> Result<()> {
21        if meta.path.is_ident("name") {
22            match meta.value() {
23                Ok(value) => {
24                    self.name = Some(value.parse::<LitStr>()?.parse()?);
25                }
26                Err(_) => {
27                    return Err(meta.error("`name` attribute must have a value"));
28                }
29            }
30        } else if meta.path.is_ident("skip_memory_check") {
31            if meta.value().is_ok() {
32                return Err(meta.error("`skip_memory_check` attribute have no values"));
33            }
34            self.skip_memory_check = true;
35        } else {
36            return Err(meta.error("unsupported module attribute"));
37        }
38        Ok(())
39    }
40}
41
42#[proc_macro_attribute]
43pub fn lua_module(attr: TokenStream, item: TokenStream) -> TokenStream {
44    let mut args = ModuleAttributes::default();
45    if !attr.is_empty() {
46        let args_parser = syn::meta::parser(|meta| args.parse(meta));
47        parse_macro_input!(attr with args_parser);
48    }
49
50    let func = parse_macro_input!(item as ItemFn);
51    let func_name = &func.sig.ident;
52    let module_name = args.name.unwrap_or_else(|| func_name.clone());
53    let ext_entrypoint_name = Ident::new(&format!("luaopen_{module_name}"), Span::call_site());
54    let skip_memory_check = if args.skip_memory_check {
55        quote! { lua.skip_memory_check(true); }
56    } else {
57        quote! {}
58    };
59
60    let wrapped = quote! {
61        mlua::require_module_feature!();
62
63        #func
64
65        #[no_mangle]
66        unsafe extern "C-unwind" fn #ext_entrypoint_name(state: *mut mlua::lua_State) -> ::std::os::raw::c_int {
67            mlua::Lua::entrypoint1(state, move |lua| {
68                #skip_memory_check
69                #func_name(lua)
70            })
71        }
72    };
73
74    wrapped.into()
75}
76
77#[cfg(feature = "macros")]
78fn to_ident(tt: &TokenTree) -> TokenStream2 {
79    let s: TokenStream = tt.clone().into();
80    s.into()
81}
82
83#[cfg(feature = "macros")]
84#[proc_macro]
85#[proc_macro_error]
86pub fn chunk(input: TokenStream) -> TokenStream {
87    let chunk = Chunk::new(input);
88
89    let source = chunk.source();
90
91    let caps_len = chunk.captures().len();
92    let caps = chunk.captures().iter().map(|cap| {
93        let cap_name = cap.as_rust().to_string();
94        let cap = to_ident(cap.as_rust());
95        quote! { env.raw_set(#cap_name, #cap)?; }
96    });
97
98    let wrapped_code = quote! {{
99        use mlua::{AsChunk, ChunkMode, Lua, Result, Table};
100        use ::std::borrow::Cow;
101        use ::std::cell::Cell;
102        use ::std::io::Result as IoResult;
103
104        struct InnerChunk<F: FnOnce(&Lua) -> Result<Table>>(Cell<Option<F>>);
105
106        impl<F> AsChunk for InnerChunk<F>
107        where
108            F: FnOnce(&Lua) -> Result<Table>,
109        {
110            fn environment(&self, lua: &Lua) -> Result<Option<Table>> {
111                if #caps_len > 0 {
112                    if let Some(make_env) = self.0.take() {
113                        return make_env(lua).map(Some);
114                    }
115                }
116                Ok(None)
117            }
118
119            fn mode(&self) -> Option<ChunkMode> {
120                Some(ChunkMode::Text)
121            }
122
123            fn source<'a>(&self) -> IoResult<Cow<'a, [u8]>> {
124                Ok(Cow::Borrowed((#source).as_bytes()))
125            }
126        }
127
128        let make_env = move |lua: &Lua| -> Result<Table> {
129            let globals = lua.globals();
130            let env = lua.create_table()?;
131            let meta = lua.create_table()?;
132            meta.raw_set("__index", &globals)?;
133            meta.raw_set("__newindex", &globals)?;
134
135            // Add captured variables
136            #(#caps)*
137
138            env.set_metatable(Some(meta))?;
139            Ok(env)
140        };
141
142        InnerChunk(Cell::new(Some(make_env)))
143    }};
144
145    wrapped_code.into()
146}
147
148#[cfg(feature = "macros")]
149#[proc_macro_derive(FromLua)]
150pub fn from_lua(input: TokenStream) -> TokenStream {
151    from_lua::from_lua(input)
152}
153
154#[cfg(feature = "macros")]
155mod chunk;
156#[cfg(feature = "macros")]
157mod from_lua;
158#[cfg(feature = "macros")]
159mod token;