mlua/
chunk.rs

1use std::borrow::Cow;
2use std::collections::HashMap;
3use std::ffi::CString;
4use std::io::Result as IoResult;
5use std::path::{Path, PathBuf};
6use std::string::String as StdString;
7
8use crate::error::{Error, ErrorContext, Result};
9use crate::function::Function;
10use crate::lua::Lua;
11use crate::table::Table;
12use crate::value::{FromLuaMulti, IntoLua, IntoLuaMulti};
13
14/// Trait for types [loadable by Lua] and convertible to a [`Chunk`]
15///
16/// [loadable by Lua]: https://www.lua.org/manual/5.4/manual.html#3.3.2
17/// [`Chunk`]: crate::Chunk
18pub trait AsChunk<'lua, 'a> {
19    /// Returns optional chunk name
20    fn name(&self) -> Option<StdString> {
21        None
22    }
23
24    /// Returns optional chunk [environment]
25    ///
26    /// [environment]: https://www.lua.org/manual/5.4/manual.html#2.2
27    fn environment(&self, lua: &'lua Lua) -> Result<Option<Table<'lua>>> {
28        let _lua = lua; // suppress warning
29        Ok(None)
30    }
31
32    /// Returns optional chunk mode (text or binary)
33    fn mode(&self) -> Option<ChunkMode> {
34        None
35    }
36
37    /// Returns chunk data (can be text or binary)
38    fn source(self) -> IoResult<Cow<'a, [u8]>>;
39}
40
41impl<'a> AsChunk<'_, 'a> for &'a str {
42    fn source(self) -> IoResult<Cow<'a, [u8]>> {
43        Ok(Cow::Borrowed(self.as_ref()))
44    }
45}
46
47impl AsChunk<'_, 'static> for StdString {
48    fn source(self) -> IoResult<Cow<'static, [u8]>> {
49        Ok(Cow::Owned(self.into_bytes()))
50    }
51}
52
53impl<'a> AsChunk<'_, 'a> for &'a StdString {
54    fn source(self) -> IoResult<Cow<'a, [u8]>> {
55        Ok(Cow::Borrowed(self.as_bytes()))
56    }
57}
58
59impl<'a> AsChunk<'_, 'a> for &'a [u8] {
60    fn source(self) -> IoResult<Cow<'a, [u8]>> {
61        Ok(Cow::Borrowed(self))
62    }
63}
64
65impl AsChunk<'_, 'static> for Vec<u8> {
66    fn source(self) -> IoResult<Cow<'static, [u8]>> {
67        Ok(Cow::Owned(self))
68    }
69}
70
71impl<'a> AsChunk<'_, 'a> for &'a Vec<u8> {
72    fn source(self) -> IoResult<Cow<'a, [u8]>> {
73        Ok(Cow::Borrowed(self.as_ref()))
74    }
75}
76
77impl AsChunk<'_, 'static> for &Path {
78    fn name(&self) -> Option<StdString> {
79        Some(format!("@{}", self.display()))
80    }
81
82    fn source(self) -> IoResult<Cow<'static, [u8]>> {
83        std::fs::read(self).map(Cow::Owned)
84    }
85}
86
87impl AsChunk<'_, 'static> for PathBuf {
88    fn name(&self) -> Option<StdString> {
89        Some(format!("@{}", self.display()))
90    }
91
92    fn source(self) -> IoResult<Cow<'static, [u8]>> {
93        std::fs::read(self).map(Cow::Owned)
94    }
95}
96
97/// Returned from [`Lua::load`] and is used to finalize loading and executing Lua main chunks.
98///
99/// [`Lua::load`]: crate::Lua::load
100#[must_use = "`Chunk`s do nothing unless one of `exec`, `eval`, `call`, or `into_function` are called on them"]
101pub struct Chunk<'lua, 'a> {
102    pub(crate) lua: &'lua Lua,
103    pub(crate) name: StdString,
104    pub(crate) env: Result<Option<Table<'lua>>>,
105    pub(crate) mode: Option<ChunkMode>,
106    pub(crate) source: IoResult<Cow<'a, [u8]>>,
107    #[cfg(feature = "luau")]
108    pub(crate) compiler: Option<Compiler>,
109}
110
111/// Represents chunk mode (text or binary).
112#[derive(Clone, Copy, Debug, PartialEq, Eq)]
113pub enum ChunkMode {
114    Text,
115    Binary,
116}
117
118/// Luau compiler
119#[cfg(any(feature = "luau", doc))]
120#[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
121#[derive(Clone, Debug)]
122pub struct Compiler {
123    optimization_level: u8,
124    debug_level: u8,
125    type_info_level: u8,
126    coverage_level: u8,
127    vector_lib: Option<String>,
128    vector_ctor: Option<String>,
129    vector_type: Option<String>,
130    mutable_globals: Vec<String>,
131    userdata_types: Vec<String>,
132}
133
134#[cfg(any(feature = "luau", doc))]
135impl Default for Compiler {
136    fn default() -> Self {
137        Self::new()
138    }
139}
140
141#[cfg(any(feature = "luau", doc))]
142impl Compiler {
143    /// Creates Luau compiler instance with default options
144    pub const fn new() -> Self {
145        // Defaults are taken from luacode.h
146        Compiler {
147            optimization_level: 1,
148            debug_level: 1,
149            type_info_level: 0,
150            coverage_level: 0,
151            vector_lib: None,
152            vector_ctor: None,
153            vector_type: None,
154            mutable_globals: Vec::new(),
155            userdata_types: Vec::new(),
156        }
157    }
158
159    /// Sets Luau compiler optimization level.
160    ///
161    /// Possible values:
162    /// * 0 - no optimization
163    /// * 1 - baseline optimization level that doesn't prevent debuggability (default)
164    /// * 2 - includes optimizations that harm debuggability such as inlining
165    #[must_use]
166    pub const fn set_optimization_level(mut self, level: u8) -> Self {
167        self.optimization_level = level;
168        self
169    }
170
171    /// Sets Luau compiler debug level.
172    ///
173    /// Possible values:
174    /// * 0 - no debugging support
175    /// * 1 - line info & function names only; sufficient for backtraces (default)
176    /// * 2 - full debug info with local & upvalue names; necessary for debugger
177    #[must_use]
178    pub const fn set_debug_level(mut self, level: u8) -> Self {
179        self.debug_level = level;
180        self
181    }
182
183    /// Sets Luau type information level used to guide native code generation decisions.
184    ///
185    /// Possible values:
186    /// * 0 - generate for native modules (default)
187    /// * 1 - generate for all modules
188    pub const fn set_type_info_level(mut self, level: u8) -> Self {
189        self.type_info_level = level;
190        self
191    }
192
193    /// Sets Luau compiler code coverage level.
194    ///
195    /// Possible values:
196    /// * 0 - no code coverage support (default)
197    /// * 1 - statement coverage
198    /// * 2 - statement and expression coverage (verbose)
199    #[must_use]
200    pub const fn set_coverage_level(mut self, level: u8) -> Self {
201        self.coverage_level = level;
202        self
203    }
204
205    #[doc(hidden)]
206    #[must_use]
207    pub fn set_vector_lib(mut self, lib: impl Into<String>) -> Self {
208        self.vector_lib = Some(lib.into());
209        self
210    }
211
212    #[doc(hidden)]
213    #[must_use]
214    pub fn set_vector_ctor(mut self, ctor: impl Into<String>) -> Self {
215        self.vector_ctor = Some(ctor.into());
216        self
217    }
218
219    #[doc(hidden)]
220    #[must_use]
221    pub fn set_vector_type(mut self, r#type: impl Into<String>) -> Self {
222        self.vector_type = Some(r#type.into());
223        self
224    }
225
226    /// Sets a list of globals that are mutable.
227    ///
228    /// It disables the import optimization for fields accessed through these.
229    #[must_use]
230    pub fn set_mutable_globals(mut self, globals: Vec<String>) -> Self {
231        self.mutable_globals = globals;
232        self
233    }
234
235    /// Sets a list of userdata types that will be included in the type information.
236    #[must_use]
237    pub fn set_userdata_types(mut self, types: Vec<String>) -> Self {
238        self.userdata_types = types;
239        self
240    }
241
242    /// Compiles the `source` into bytecode.
243    pub fn compile(&self, source: impl AsRef<[u8]>) -> Vec<u8> {
244        use std::os::raw::c_int;
245        use std::ptr;
246
247        let vector_lib = self.vector_lib.clone();
248        let vector_lib = vector_lib.and_then(|lib| CString::new(lib).ok());
249        let vector_lib = vector_lib.as_ref();
250        let vector_ctor = self.vector_ctor.clone();
251        let vector_ctor = vector_ctor.and_then(|ctor| CString::new(ctor).ok());
252        let vector_ctor = vector_ctor.as_ref();
253        let vector_type = self.vector_type.clone();
254        let vector_type = vector_type.and_then(|t| CString::new(t).ok());
255        let vector_type = vector_type.as_ref();
256
257        macro_rules! vec2cstring_ptr {
258            ($name:ident, $name_ptr:ident) => {
259                let $name = self
260                    .$name
261                    .iter()
262                    .map(|name| CString::new(name.clone()).ok())
263                    .collect::<Option<Vec<_>>>()
264                    .unwrap_or_default();
265                let mut $name = $name.iter().map(|s| s.as_ptr()).collect::<Vec<_>>();
266                let mut $name_ptr = ptr::null();
267                if !$name.is_empty() {
268                    $name.push(ptr::null());
269                    $name_ptr = $name.as_ptr();
270                }
271            };
272        }
273
274        vec2cstring_ptr!(mutable_globals, mutable_globals_ptr);
275        vec2cstring_ptr!(userdata_types, userdata_types_ptr);
276
277        unsafe {
278            let mut options = ffi::lua_CompileOptions::default();
279            options.optimizationLevel = self.optimization_level as c_int;
280            options.debugLevel = self.debug_level as c_int;
281            options.typeInfoLevel = self.type_info_level as c_int;
282            options.coverageLevel = self.coverage_level as c_int;
283            options.vectorLib = vector_lib.map_or(ptr::null(), |s| s.as_ptr());
284            options.vectorCtor = vector_ctor.map_or(ptr::null(), |s| s.as_ptr());
285            options.vectorType = vector_type.map_or(ptr::null(), |s| s.as_ptr());
286            options.mutableGlobals = mutable_globals_ptr;
287            options.userdataTypes = userdata_types_ptr;
288            ffi::luau_compile(source.as_ref(), options)
289        }
290    }
291}
292
293impl<'lua, 'a> Chunk<'lua, 'a> {
294    /// Sets the name of this chunk, which results in more informative error traces.
295    pub fn set_name(mut self, name: impl Into<String>) -> Self {
296        self.name = name.into();
297        self
298    }
299
300    /// Sets the environment of the loaded chunk to the given value.
301    ///
302    /// In Lua >=5.2 main chunks always have exactly one upvalue, and this upvalue is used as the `_ENV`
303    /// variable inside the chunk. By default this value is set to the global environment.
304    ///
305    /// Calling this method changes the `_ENV` upvalue to the value provided, and variables inside
306    /// the chunk will refer to the given environment rather than the global one.
307    ///
308    /// All global variables (including the standard library!) are looked up in `_ENV`, so it may be
309    /// necessary to populate the environment in order for scripts using custom environments to be
310    /// useful.
311    pub fn set_environment<V: IntoLua<'lua>>(mut self, env: V) -> Self {
312        self.env = env
313            .into_lua(self.lua)
314            .and_then(|val| self.lua.unpack(val))
315            .context("bad environment value");
316        self
317    }
318
319    /// Sets whether the chunk is text or binary (autodetected by default).
320    ///
321    /// Be aware, Lua does not check the consistency of the code inside binary chunks.
322    /// Running maliciously crafted bytecode can crash the interpreter.
323    pub fn set_mode(mut self, mode: ChunkMode) -> Self {
324        self.mode = Some(mode);
325        self
326    }
327
328    /// Sets or overwrites a Luau compiler used for this chunk.
329    ///
330    /// See [`Compiler`] for details and possible options.
331    ///
332    /// Requires `feature = "luau"`
333    #[cfg(any(feature = "luau", doc))]
334    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
335    pub fn set_compiler(mut self, compiler: Compiler) -> Self {
336        self.compiler = Some(compiler);
337        self
338    }
339
340    /// Execute this chunk of code.
341    ///
342    /// This is equivalent to calling the chunk function with no arguments and no return values.
343    pub fn exec(self) -> Result<()> {
344        self.call::<_, ()>(())?;
345        Ok(())
346    }
347
348    /// Asynchronously execute this chunk of code.
349    ///
350    /// See [`exec`] for more details.
351    ///
352    /// Requires `feature = "async"`
353    ///
354    /// [`exec`]: #method.exec
355    #[cfg(feature = "async")]
356    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
357    pub async fn exec_async(self) -> Result<()> {
358        self.call_async(()).await
359    }
360
361    /// Evaluate the chunk as either an expression or block.
362    ///
363    /// If the chunk can be parsed as an expression, this loads and executes the chunk and returns
364    /// the value that it evaluates to. Otherwise, the chunk is interpreted as a block as normal,
365    /// and this is equivalent to calling `exec`.
366    pub fn eval<R: FromLuaMulti<'lua>>(self) -> Result<R> {
367        // Bytecode is always interpreted as a statement.
368        // For source code, first try interpreting the lua as an expression by adding
369        // "return", then as a statement. This is the same thing the
370        // actual lua repl does.
371        if self.detect_mode() == ChunkMode::Binary {
372            self.call(())
373        } else if let Ok(function) = self.to_expression() {
374            function.call(())
375        } else {
376            self.call(())
377        }
378    }
379
380    /// Asynchronously evaluate the chunk as either an expression or block.
381    ///
382    /// See [`eval`] for more details.
383    ///
384    /// Requires `feature = "async"`
385    ///
386    /// [`eval`]: #method.eval
387    #[cfg(feature = "async")]
388    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
389    pub async fn eval_async<R>(self) -> Result<R>
390    where
391        R: FromLuaMulti<'lua> + 'lua,
392    {
393        if self.detect_mode() == ChunkMode::Binary {
394            self.call_async(()).await
395        } else if let Ok(function) = self.to_expression() {
396            function.call_async(()).await
397        } else {
398            self.call_async(()).await
399        }
400    }
401
402    /// Load the chunk function and call it with the given arguments.
403    ///
404    /// This is equivalent to `into_function` and calling the resulting function.
405    pub fn call<A: IntoLuaMulti<'lua>, R: FromLuaMulti<'lua>>(self, args: A) -> Result<R> {
406        self.into_function()?.call(args)
407    }
408
409    /// Load the chunk function and asynchronously call it with the given arguments.
410    ///
411    /// See [`call`] for more details.
412    ///
413    /// Requires `feature = "async"`
414    ///
415    /// [`call`]: #method.call
416    #[cfg(feature = "async")]
417    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
418    pub async fn call_async<A, R>(self, args: A) -> Result<R>
419    where
420        A: IntoLuaMulti<'lua>,
421        R: FromLuaMulti<'lua> + 'lua,
422    {
423        self.into_function()?.call_async(args).await
424    }
425
426    /// Load this chunk into a regular `Function`.
427    ///
428    /// This simply compiles the chunk without actually executing it.
429    #[cfg_attr(not(feature = "luau"), allow(unused_mut))]
430    pub fn into_function(mut self) -> Result<Function<'lua>> {
431        #[cfg(feature = "luau")]
432        if self.compiler.is_some() {
433            // We don't need to compile source if no compiler set
434            self.compile();
435        }
436
437        let name = Self::convert_name(self.name)?;
438        self.lua
439            .load_chunk(Some(&name), self.env?, self.mode, self.source?.as_ref())
440    }
441
442    /// Compiles the chunk and changes mode to binary.
443    ///
444    /// It does nothing if the chunk is already binary.
445    fn compile(&mut self) {
446        if let Ok(ref source) = self.source {
447            if self.detect_mode() == ChunkMode::Text {
448                #[cfg(feature = "luau")]
449                {
450                    let data = self
451                        .compiler
452                        .get_or_insert_with(Default::default)
453                        .compile(source);
454                    self.source = Ok(Cow::Owned(data));
455                    self.mode = Some(ChunkMode::Binary);
456                }
457                #[cfg(not(feature = "luau"))]
458                if let Ok(func) = self.lua.load_chunk(None, None, None, source.as_ref()) {
459                    let data = func.dump(false);
460                    self.source = Ok(Cow::Owned(data));
461                    self.mode = Some(ChunkMode::Binary);
462                }
463            }
464        }
465    }
466
467    /// Fetches compiled bytecode of this chunk from the cache.
468    ///
469    /// If not found, compiles the source code and stores it on the cache.
470    pub(crate) fn try_cache(mut self) -> Self {
471        struct ChunksCache(HashMap<Vec<u8>, Vec<u8>>);
472
473        // Try to fetch compiled chunk from cache
474        let mut text_source = None;
475        if let Ok(ref source) = self.source {
476            if self.detect_mode() == ChunkMode::Text {
477                if let Some(cache) = self.lua.app_data_ref::<ChunksCache>() {
478                    if let Some(data) = cache.0.get(source.as_ref()) {
479                        self.source = Ok(Cow::Owned(data.clone()));
480                        self.mode = Some(ChunkMode::Binary);
481                        return self;
482                    }
483                }
484                text_source = Some(source.as_ref().to_vec());
485            }
486        }
487
488        // Compile and cache the chunk
489        if let Some(text_source) = text_source {
490            self.compile();
491            if let Ok(ref binary_source) = self.source {
492                if self.detect_mode() == ChunkMode::Binary {
493                    if let Some(mut cache) = self.lua.app_data_mut::<ChunksCache>() {
494                        cache.0.insert(text_source, binary_source.as_ref().to_vec());
495                    } else {
496                        let mut cache = ChunksCache(HashMap::new());
497                        cache.0.insert(text_source, binary_source.as_ref().to_vec());
498                        let _ = self.lua.try_set_app_data(cache);
499                    }
500                }
501            }
502        }
503
504        self
505    }
506
507    fn to_expression(&self) -> Result<Function<'lua>> {
508        // We assume that mode is Text
509        let source = self.source.as_ref();
510        let source = source.map_err(Error::runtime)?;
511        let source = Self::expression_source(source);
512        // We don't need to compile source if no compiler options set
513        #[cfg(feature = "luau")]
514        let source = self
515            .compiler
516            .as_ref()
517            .map(|c| c.compile(&source))
518            .unwrap_or(source);
519
520        let name = Self::convert_name(self.name.clone())?;
521        self.lua
522            .load_chunk(Some(&name), self.env.clone()?, None, &source)
523    }
524
525    fn detect_mode(&self) -> ChunkMode {
526        match (self.mode, &self.source) {
527            (Some(mode), _) => mode,
528            (None, Ok(source)) => {
529                #[cfg(not(feature = "luau"))]
530                if source.starts_with(ffi::LUA_SIGNATURE) {
531                    return ChunkMode::Binary;
532                }
533                #[cfg(feature = "luau")]
534                if *source.first().unwrap_or(&u8::MAX) < b'\n' {
535                    return ChunkMode::Binary;
536                }
537                ChunkMode::Text
538            }
539            (None, Err(_)) => ChunkMode::Text, // any value is fine
540        }
541    }
542
543    fn convert_name(name: String) -> Result<CString> {
544        CString::new(name).map_err(|err| Error::runtime(format!("invalid name: {err}")))
545    }
546
547    fn expression_source(source: &[u8]) -> Vec<u8> {
548        let mut buf = Vec::with_capacity(b"return ".len() + source.len());
549        buf.extend(b"return ");
550        buf.extend(source);
551        buf
552    }
553}