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 #(#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;