mlua_codemp_patch/
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, Result};
9use crate::function::Function;
10use crate::state::{Lua, WeakLua};
11use crate::table::Table;
12use crate::value::{FromLuaMulti, 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<'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) -> Result<Option<Table>> {
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<'a> {
102    pub(crate) lua: WeakLua,
103    pub(crate) name: StdString,
104    pub(crate) env: Result<Option<Table>>,
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        const { 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    ///
244    /// Returns `Error::SyntaxError` if the source code is invalid.
245    pub fn compile(&self, source: impl AsRef<[u8]>) -> Result<Vec<u8>> {
246        use std::os::raw::c_int;
247        use std::ptr;
248
249        let vector_lib = self.vector_lib.clone();
250        let vector_lib = vector_lib.and_then(|lib| CString::new(lib).ok());
251        let vector_lib = vector_lib.as_ref();
252        let vector_ctor = self.vector_ctor.clone();
253        let vector_ctor = vector_ctor.and_then(|ctor| CString::new(ctor).ok());
254        let vector_ctor = vector_ctor.as_ref();
255        let vector_type = self.vector_type.clone();
256        let vector_type = vector_type.and_then(|t| CString::new(t).ok());
257        let vector_type = vector_type.as_ref();
258
259        macro_rules! vec2cstring_ptr {
260            ($name:ident, $name_ptr:ident) => {
261                let $name = self
262                    .$name
263                    .iter()
264                    .map(|name| CString::new(name.clone()).ok())
265                    .collect::<Option<Vec<_>>>()
266                    .unwrap_or_default();
267                let mut $name = $name.iter().map(|s| s.as_ptr()).collect::<Vec<_>>();
268                let mut $name_ptr = ptr::null();
269                if !$name.is_empty() {
270                    $name.push(ptr::null());
271                    $name_ptr = $name.as_ptr();
272                }
273            };
274        }
275
276        vec2cstring_ptr!(mutable_globals, mutable_globals_ptr);
277        vec2cstring_ptr!(userdata_types, userdata_types_ptr);
278
279        let bytecode = unsafe {
280            let mut options = ffi::lua_CompileOptions::default();
281            options.optimizationLevel = self.optimization_level as c_int;
282            options.debugLevel = self.debug_level as c_int;
283            options.typeInfoLevel = self.type_info_level as c_int;
284            options.coverageLevel = self.coverage_level as c_int;
285            options.vectorLib = vector_lib.map_or(ptr::null(), |s| s.as_ptr());
286            options.vectorCtor = vector_ctor.map_or(ptr::null(), |s| s.as_ptr());
287            options.vectorType = vector_type.map_or(ptr::null(), |s| s.as_ptr());
288            options.mutableGlobals = mutable_globals_ptr;
289            options.userdataTypes = userdata_types_ptr;
290            ffi::luau_compile(source.as_ref(), options)
291        };
292
293        if bytecode.first() == Some(&0) {
294            // The rest of the bytecode is the error message starting with `:`
295            // See https://github.com/luau-lang/luau/blob/0.640/Compiler/src/Compiler.cpp#L4336
296            let message = String::from_utf8_lossy(&bytecode[2..]).to_string();
297            return Err(Error::SyntaxError {
298                incomplete_input: message.ends_with("<eof>"),
299                message,
300            });
301        }
302
303        Ok(bytecode)
304    }
305}
306
307impl<'a> Chunk<'a> {
308    /// Sets the name of this chunk, which results in more informative error traces.
309    pub fn set_name(mut self, name: impl Into<String>) -> Self {
310        self.name = name.into();
311        self
312    }
313
314    /// Sets the environment of the loaded chunk to the given value.
315    ///
316    /// In Lua >=5.2 main chunks always have exactly one upvalue, and this upvalue is used as the
317    /// `_ENV` variable inside the chunk. By default this value is set to the global environment.
318    ///
319    /// Calling this method changes the `_ENV` upvalue to the value provided, and variables inside
320    /// the chunk will refer to the given environment rather than the global one.
321    ///
322    /// All global variables (including the standard library!) are looked up in `_ENV`, so it may be
323    /// necessary to populate the environment in order for scripts using custom environments to be
324    /// useful.
325    pub fn set_environment(mut self, env: Table) -> Self {
326        self.env = Ok(Some(env));
327        self
328    }
329
330    /// Sets whether the chunk is text or binary (autodetected by default).
331    ///
332    /// Be aware, Lua does not check the consistency of the code inside binary chunks.
333    /// Running maliciously crafted bytecode can crash the interpreter.
334    pub fn set_mode(mut self, mode: ChunkMode) -> Self {
335        self.mode = Some(mode);
336        self
337    }
338
339    /// Sets or overwrites a Luau compiler used for this chunk.
340    ///
341    /// See [`Compiler`] for details and possible options.
342    ///
343    /// Requires `feature = "luau"`
344    #[cfg(any(feature = "luau", doc))]
345    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
346    pub fn set_compiler(mut self, compiler: Compiler) -> Self {
347        self.compiler = Some(compiler);
348        self
349    }
350
351    /// Execute this chunk of code.
352    ///
353    /// This is equivalent to calling the chunk function with no arguments and no return values.
354    pub fn exec(self) -> Result<()> {
355        self.call(())
356    }
357
358    /// Asynchronously execute this chunk of code.
359    ///
360    /// See [`exec`] for more details.
361    ///
362    /// Requires `feature = "async"`
363    ///
364    /// [`exec`]: #method.exec
365    #[cfg(feature = "async")]
366    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
367    pub async fn exec_async(self) -> Result<()> {
368        self.call_async(()).await
369    }
370
371    /// Evaluate the chunk as either an expression or block.
372    ///
373    /// If the chunk can be parsed as an expression, this loads and executes the chunk and returns
374    /// the value that it evaluates to. Otherwise, the chunk is interpreted as a block as normal,
375    /// and this is equivalent to calling `exec`.
376    pub fn eval<R: FromLuaMulti>(self) -> Result<R> {
377        // Bytecode is always interpreted as a statement.
378        // For source code, first try interpreting the lua as an expression by adding
379        // "return", then as a statement. This is the same thing the
380        // actual lua repl does.
381        if self.detect_mode() == ChunkMode::Binary {
382            self.call(())
383        } else if let Ok(function) = self.to_expression() {
384            function.call(())
385        } else {
386            self.call(())
387        }
388    }
389
390    /// Asynchronously evaluate the chunk as either an expression or block.
391    ///
392    /// See [`eval`] for more details.
393    ///
394    /// Requires `feature = "async"`
395    ///
396    /// [`eval`]: #method.eval
397    #[cfg(feature = "async")]
398    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
399    pub async fn eval_async<R>(self) -> Result<R>
400    where
401        R: FromLuaMulti,
402    {
403        if self.detect_mode() == ChunkMode::Binary {
404            self.call_async(()).await
405        } else if let Ok(function) = self.to_expression() {
406            function.call_async(()).await
407        } else {
408            self.call_async(()).await
409        }
410    }
411
412    /// Load the chunk function and call it with the given arguments.
413    ///
414    /// This is equivalent to `into_function` and calling the resulting function.
415    pub fn call<R: FromLuaMulti>(self, args: impl IntoLuaMulti) -> Result<R> {
416        self.into_function()?.call(args)
417    }
418
419    /// Load the chunk function and asynchronously call it with the given arguments.
420    ///
421    /// See [`call`] for more details.
422    ///
423    /// Requires `feature = "async"`
424    ///
425    /// [`call`]: #method.call
426    #[cfg(feature = "async")]
427    #[cfg_attr(docsrs, doc(cfg(feature = "async")))]
428    pub async fn call_async<R>(self, args: impl IntoLuaMulti) -> Result<R>
429    where
430        R: FromLuaMulti,
431    {
432        self.into_function()?.call_async(args).await
433    }
434
435    /// Load this chunk into a regular `Function`.
436    ///
437    /// This simply compiles the chunk without actually executing it.
438    #[cfg_attr(not(feature = "luau"), allow(unused_mut))]
439    pub fn into_function(mut self) -> Result<Function> {
440        #[cfg(feature = "luau")]
441        if self.compiler.is_some() {
442            // We don't need to compile source if no compiler set
443            self.compile();
444        }
445
446        let name = Self::convert_name(self.name)?;
447        self.lua
448            .lock()
449            .load_chunk(Some(&name), self.env?.as_ref(), self.mode, self.source?.as_ref())
450    }
451
452    /// Compiles the chunk and changes mode to binary.
453    ///
454    /// It does nothing if the chunk is already binary or invalid.
455    fn compile(&mut self) {
456        if let Ok(ref source) = self.source {
457            if self.detect_mode() == ChunkMode::Text {
458                #[cfg(feature = "luau")]
459                if let Ok(data) = self.compiler.get_or_insert_with(Default::default).compile(source) {
460                    self.source = Ok(Cow::Owned(data));
461                    self.mode = Some(ChunkMode::Binary);
462                }
463                #[cfg(not(feature = "luau"))]
464                if let Ok(func) = self.lua.lock().load_chunk(None, None, None, source.as_ref()) {
465                    let data = func.dump(false);
466                    self.source = Ok(Cow::Owned(data));
467                    self.mode = Some(ChunkMode::Binary);
468                }
469            }
470        }
471    }
472
473    /// Fetches compiled bytecode of this chunk from the cache.
474    ///
475    /// If not found, compiles the source code and stores it on the cache.
476    pub(crate) fn try_cache(mut self) -> Self {
477        struct ChunksCache(HashMap<Vec<u8>, Vec<u8>>);
478
479        // Try to fetch compiled chunk from cache
480        let mut text_source = None;
481        if let Ok(ref source) = self.source {
482            if self.detect_mode() == ChunkMode::Text {
483                let lua = self.lua.lock();
484                if let Some(cache) = lua.app_data_ref::<ChunksCache>() {
485                    if let Some(data) = cache.0.get(source.as_ref()) {
486                        self.source = Ok(Cow::Owned(data.clone()));
487                        self.mode = Some(ChunkMode::Binary);
488                        return self;
489                    }
490                }
491                text_source = Some(source.as_ref().to_vec());
492            }
493        }
494
495        // Compile and cache the chunk
496        if let Some(text_source) = text_source {
497            self.compile();
498            if let Ok(ref binary_source) = self.source {
499                if self.detect_mode() == ChunkMode::Binary {
500                    let lua = self.lua.lock();
501                    if let Some(mut cache) = lua.app_data_mut::<ChunksCache>() {
502                        cache.0.insert(text_source, binary_source.as_ref().to_vec());
503                    } else {
504                        let mut cache = ChunksCache(HashMap::new());
505                        cache.0.insert(text_source, binary_source.as_ref().to_vec());
506                        let _ = lua.try_set_app_data(cache);
507                    };
508                }
509            }
510        }
511
512        self
513    }
514
515    fn to_expression(&self) -> Result<Function> {
516        // We assume that mode is Text
517        let source = self.source.as_ref();
518        let source = source.map_err(Error::runtime)?;
519        let source = Self::expression_source(source);
520        // We don't need to compile source if no compiler options set
521        #[cfg(feature = "luau")]
522        let source = self
523            .compiler
524            .as_ref()
525            .map(|c| c.compile(&source))
526            .transpose()?
527            .unwrap_or(source);
528
529        let name = Self::convert_name(self.name.clone())?;
530        let env = match &self.env {
531            Ok(Some(env)) => Some(env),
532            Ok(None) => None,
533            Err(err) => return Err(err.clone()),
534        };
535        self.lua.lock().load_chunk(Some(&name), env, None, &source)
536    }
537
538    fn detect_mode(&self) -> ChunkMode {
539        match (self.mode, &self.source) {
540            (Some(mode), _) => mode,
541            (None, Ok(source)) => {
542                #[cfg(not(feature = "luau"))]
543                if source.starts_with(ffi::LUA_SIGNATURE) {
544                    return ChunkMode::Binary;
545                }
546                #[cfg(feature = "luau")]
547                if *source.first().unwrap_or(&u8::MAX) < b'\n' {
548                    return ChunkMode::Binary;
549                }
550                ChunkMode::Text
551            }
552            (None, Err(_)) => ChunkMode::Text, // any value is fine
553        }
554    }
555
556    fn convert_name(name: String) -> Result<CString> {
557        CString::new(name).map_err(|err| Error::runtime(format!("invalid name: {err}")))
558    }
559
560    fn expression_source(source: &[u8]) -> Vec<u8> {
561        let mut buf = Vec::with_capacity(b"return ".len() + source.len());
562        buf.extend(b"return ");
563        buf.extend(source);
564        buf
565    }
566}
mlua_codemp_patch/chunk.rs

mlua_codemp_patch/
chunk.rs
