luau 0.732.0

use std::borrow::Cow;
use std::io;
use std::panic::Location;
use std::path::{Path, PathBuf};

use luau_bytecode::opcodes::{
    BYTECODE_TYPE_VERSION_MAX, BYTECODE_TYPE_VERSION_MIN, BYTECODE_VERSION_CLASSES,
    BYTECODE_VERSION_MAX,
};

use crate::error::Error;
use crate::function::Function;
use crate::lua::Compiler;
use crate::lua::{ChunkLoad, Lua, LuaRef};
use crate::table::Table;
use crate::thread::Thread;
use crate::value::Value;
use crate::value::{FromLuaMulti, IntoLua, IntoLuaMulti};

#[cfg(feature = "macros")]
mod capture;

#[cfg(feature = "macros")]
#[doc(hidden)]
pub use capture::{CaptureEnvironment, CapturedChunk, captured_chunk};

/// A source of Luau code or bytecode accepted by [`Lua::load`].
pub trait AsChunk<'lua> {
    /// Returns the optional chunk name.
    fn name(&self) -> Option<String> {
        None
    }

    /// Returns the optional chunk environment.
    fn environment(&self, lua: &LuaRef<'lua>) -> Result<Option<Table<'lua>>, Error> {
        let _ = lua;
        Ok(None)
    }

    /// Returns the optional text or binary mode.
    fn mode(&self) -> Option<ChunkMode> {
        None
    }

    /// Returns the chunk contents.
    fn source(&self) -> io::Result<Cow<'lua, [u8]>>;
}

impl<'lua> AsChunk<'lua> for &'lua str {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Borrowed(self.as_bytes()))
    }
}

impl<'lua> AsChunk<'lua> for String {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Owned(self.as_bytes().to_vec()))
    }
}

impl<'lua> AsChunk<'lua> for &'lua String {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Borrowed(self.as_bytes()))
    }
}

impl<'lua> AsChunk<'lua> for Cow<'lua, str> {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(match self {
            Cow::Borrowed(source) => Cow::Borrowed(source.as_bytes()),
            Cow::Owned(source) => Cow::Owned(source.as_bytes().to_vec()),
        })
    }
}

impl<'lua> AsChunk<'lua> for &'lua [u8] {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Borrowed(self))
    }
}

impl<'lua, const N: usize> AsChunk<'lua> for &'lua [u8; N] {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Borrowed(&self[..]))
    }
}

impl<'lua> AsChunk<'lua> for Vec<u8> {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Owned(self.clone()))
    }
}

impl<'lua> AsChunk<'lua> for &'lua Vec<u8> {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(Cow::Borrowed(self.as_slice()))
    }
}

impl<'lua> AsChunk<'lua> for Cow<'lua, [u8]> {
    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        Ok(match self {
            Cow::Borrowed(source) => Cow::Borrowed(source),
            Cow::Owned(source) => Cow::Owned(source.clone()),
        })
    }
}

impl<'lua> AsChunk<'lua> for &'lua Path {
    fn name(&self) -> Option<String> {
        Some(format!("@{}", self.display()))
    }

    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        std::fs::read(self).map(Cow::Owned)
    }
}

impl<'lua> AsChunk<'lua> for PathBuf {
    fn name(&self) -> Option<String> {
        Some(format!("@{}", self.display()))
    }

    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        std::fs::read(self).map(Cow::Owned)
    }
}

impl<'lua, C> AsChunk<'lua> for Box<C>
where
    C: AsChunk<'lua> + ?Sized,
{
    fn name(&self) -> Option<String> {
        (**self).name()
    }

    fn environment(&self, lua: &LuaRef<'lua>) -> Result<Option<Table<'lua>>, Error> {
        (**self).environment(lua)
    }

    fn mode(&self) -> Option<ChunkMode> {
        (**self).mode()
    }

    fn source(&self) -> io::Result<Cow<'lua, [u8]>> {
        (**self).source()
    }
}

/// The representation of a loaded chunk.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ChunkMode {
    /// Luau source text.
    Text,
    /// Luau bytecode.
    Binary,
}

/// A configurable Luau chunk returned by [`Lua::load`].
#[must_use = "`Chunk`s do nothing unless one of `exec`, `eval`, `call`, `into_function`, `into_thread`, or `into_sandboxed` is called on them"]
pub struct Chunk<'lua> {
    thread: Thread<'lua>,
    load: ChunkLoad,
    source: io::Result<Cow<'lua, [u8]>>,
    name: String,
    environment: Result<Option<Table<'lua>>, Error>,
    mode: Option<ChunkMode>,
    compiler: Compiler,
}

/// A chunk loaded onto a private sandboxed thread but not yet executed.
///
/// The loaded function can be inspected without adding a second registry
/// root. Consuming this value yields the thread that owns the function.
pub struct SandboxedChunk<'lua> {
    thread: Thread<'lua>,
}

impl Lua {
    /// Returns a builder for loading Luau source code or bytecode.
    ///
    /// The source is not compiled or executed until the returned [`Chunk`] is
    /// consumed.
    #[track_caller]
    pub fn load<'lua>(&'lua self, source: impl AsChunk<'lua>) -> Chunk<'lua> {
        let lua = self.lua_ref();
        Chunk::new(
            lua.current_thread(),
            ChunkLoad::Main,
            &lua,
            source,
            self.runtime.compiler(),
            location_chunk_name(Location::caller()),
        )
    }
}

impl<'lua> LuaRef<'lua> {
    /// Returns a builder for loading Luau source code or bytecode.
    ///
    /// The source is not compiled or executed until the returned [`Chunk`] is
    /// consumed.
    #[track_caller]
    pub fn load(&self, source: impl AsChunk<'lua>) -> Chunk<'lua> {
        self.load_with_location(source, Location::caller())
    }

    pub(crate) fn load_with_location(
        &self,
        source: impl AsChunk<'lua>,
        location: &'static Location<'static>,
    ) -> Chunk<'lua> {
        Chunk::new(
            self.current_thread(),
            ChunkLoad::Dynamic,
            self,
            source,
            self.runtime().compiler(),
            location_chunk_name(location),
        )
    }
}

impl<'lua> Chunk<'lua> {
    pub(crate) fn new(
        thread: Thread<'lua>,
        load: ChunkLoad,
        lua: &LuaRef<'lua>,
        source: impl AsChunk<'lua>,
        compiler: Compiler,
        default_name: String,
    ) -> Self {
        let name = source.name().unwrap_or(default_name);
        let environment = source.environment(lua);
        let mode = source.mode();
        let source = source.source();
        Self {
            thread,
            load,
            source,
            name,
            environment,
            mode,
            compiler,
        }
    }

    /// Returns the name used in errors and debug information.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// Sets the name used in errors and debug information.
    ///
    /// A name beginning with `@` is treated as a file path. A name beginning
    /// with `=` is displayed without the prefix.
    pub fn set_name(mut self, name: impl Into<String>) -> Self {
        self.name = name.into();
        self
    }

    /// Returns the configured environment, if any.
    pub fn environment(&self) -> Option<&Table<'lua>> {
        self.environment.as_ref().ok()?.as_ref()
    }

    /// Uses an explicit environment for this chunk.
    ///
    /// Import constants are resolved from this environment. A custom
    /// environment's safe-environment setting remains under the embedder's
    /// control; passing the current globals participates in managed sandbox
    /// load tracking.
    pub fn set_environment(mut self, environment: Table<'lua>) -> Self {
        self.environment = Ok(Some(environment));
        self
    }

    /// Returns the configured or detected chunk mode.
    pub fn mode(&self) -> ChunkMode {
        self.detect_mode()
    }

    /// Sets whether this chunk contains source text or bytecode.
    pub fn set_mode(mut self, mode: ChunkMode) -> Self {
        self.mode = Some(mode);
        self
    }

    /// Sets the compiler used for this chunk.
    pub fn set_compiler(mut self, compiler: Compiler) -> Self {
        self.compiler = compiler;
        self
    }

    /// Compiles and loads this chunk as a function without executing it.
    pub fn into_function(self) -> Result<Function<'lua>, Error> {
        let name = self.chunk_name();
        let bytecode = self.bytecode()?;
        self.thread.load_bytecode(
            name,
            bytecode.as_ref(),
            self.environment_result()?,
            self.load,
        )
    }

    /// Loads this chunk as the body of a new coroutine.
    pub fn into_thread(self) -> Result<Thread<'lua>, Error> {
        let name = self.chunk_name();
        let bytecode = self.bytecode()?;
        let thread = self.thread.create_empty_thread()?;
        thread.load_body(
            name,
            bytecode.as_ref(),
            self.environment_result()?,
            self.load,
        )?;
        Ok(thread)
    }

    /// Loads this chunk onto a private sandboxed coroutine.
    ///
    /// Luau associates a function with the active environment when it is
    /// loaded, so sandboxing after [`Chunk::into_thread`] would be too late to
    /// isolate a chunk without an explicitly configured environment.
    /// [`Lua::sandbox(true)`](Lua::sandbox) must be enabled first.
    pub fn into_sandboxed(self) -> Result<SandboxedChunk<'lua>, Error> {
        let name = self.chunk_name();
        let bytecode = self.bytecode()?;
        let thread = self.thread.create_empty_thread()?;
        thread.sandbox_with_immutable_base()?;
        thread.load_body(
            name,
            bytecode.as_ref(),
            self.environment_result()?,
            ChunkLoad::FreshSandbox,
        )?;
        Ok(SandboxedChunk { thread })
    }

    /// Executes this chunk without arguments or return values.
    pub fn exec(self) -> Result<(), Error> {
        self.call(())
    }

    /// Executes this chunk with the given arguments.
    pub fn call<R>(self, args: impl IntoLuaMulti<'lua>) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        let name = self.chunk_name();
        let bytecode = self.bytecode()?;
        self.thread.run_bytecode_with_args(
            name,
            bytecode.as_ref(),
            args,
            self.environment_result()?,
            self.load,
        )
    }

    /// Evaluates this chunk as an expression or, if that fails, as a block.
    pub fn eval<R>(self) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        if self.mode() == ChunkMode::Binary {
            self.call(())
        } else if let Ok(function) = self.to_expression() {
            function.call(())
        } else {
            self.call(())
        }
    }

    fn to_expression(&self) -> Result<Function<'lua>, Error> {
        let source = expression_source(self.source_bytes()?);
        let bytecode = self.compile_source(source)?;
        self.thread.load_bytecode(
            self.chunk_name(),
            bytecode,
            self.environment_result()?,
            self.load,
        )
    }

    fn bytecode(&self) -> Result<Cow<'_, [u8]>, Error> {
        let source = self.source_bytes()?;
        match self.mode() {
            ChunkMode::Text => Ok(Cow::Owned(self.compile_source(source)?)),
            ChunkMode::Binary => Ok(Cow::Borrowed(source)),
        }
    }

    fn compile_source(&self, source: impl AsRef<[u8]>) -> Result<Vec<u8>, Error> {
        self.compiler.compile(source).map_err(|error| {
            let (location, detail) = match &error {
                luau_compiler::CompilerError::Parse(errors) => {
                    (errors.first().location, errors.first().to_string())
                }
                luau_compiler::CompilerError::Compile(error) => {
                    (error.location(), error.to_string())
                }
            };
            let name = self
                .name
                .strip_prefix('@')
                .or_else(|| self.name.strip_prefix('='))
                .unwrap_or(self.name.as_str());
            let message = format!("{name}:{}: {detail}", location.begin.line.saturating_add(1));
            Error::SyntaxError {
                incomplete_input: detail.ends_with("<eof>"),
                message,
            }
        })
    }

    fn chunk_name(&self) -> Vec<u8> {
        normalized_chunk_name(&self.name)
    }

    fn environment_result(&self) -> Result<Option<&Table<'lua>>, Error> {
        self.environment
            .as_ref()
            .map(|environment| environment.as_ref())
            .map_err(Clone::clone)
    }

    fn source_bytes(&self) -> Result<&[u8], Error> {
        self.source
            .as_deref()
            .map_err(|error| Error::runtime(format_args!("failed to read chunk source: {error}")))
    }

    fn detect_mode(&self) -> ChunkMode {
        if let Some(mode) = self.mode {
            return mode;
        }

        if let Ok(source) = self.source.as_deref()
            && is_luau_bytecode(source)
        {
            return ChunkMode::Binary;
        }

        ChunkMode::Text
    }
}

impl<'lua> SandboxedChunk<'lua> {
    /// Borrows the loaded entry function.
    pub fn function(&self) -> Function<'_> {
        // The private thread is loaded with exactly one function and cannot be
        // mutated without consuming this owner.
        unsafe { Function::from_borrowed_stack(&self.thread, 1) }
    }

    /// Consumes this owner and returns the prepared coroutine.
    pub fn into_thread(self) -> Thread<'lua> {
        self.thread
    }
}

struct WrappedChunk<C> {
    chunk: C,
    caller: &'static Location<'static>,
}

impl Chunk<'_> {
    /// Wraps a chunk so it can be converted into a Luau function.
    #[track_caller]
    pub fn wrap<'lua, C>(chunk: C) -> impl IntoLua<'lua>
    where
        C: AsChunk<'lua>,
    {
        WrappedChunk {
            chunk,
            caller: Location::caller(),
        }
    }
}

impl<'lua, C> IntoLua<'lua> for WrappedChunk<C>
where
    C: AsChunk<'lua>,
{
    fn into_lua(self, lua: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        let runtime = lua.runtime();
        let function = Chunk::new(
            lua.current_thread(),
            ChunkLoad::Dynamic,
            &lua,
            self.chunk,
            runtime.compiler(),
            location_chunk_name(self.caller),
        )
        .into_function()?;
        Ok(Value::Function(function))
    }
}

pub(crate) fn location_chunk_name(location: &Location<'_>) -> String {
    format!("@{}:{}", location.file(), location.line())
}

fn normalized_chunk_name(name: &str) -> Vec<u8> {
    if name.starts_with('@') || name.starts_with('=') {
        name.as_bytes().to_vec()
    } else {
        format!("={name}").into_bytes()
    }
}

fn expression_source(source: &[u8]) -> Vec<u8> {
    let mut expression = Vec::with_capacity(b"return ".len() + source.len());
    expression.extend_from_slice(b"return ");
    expression.extend_from_slice(source);
    expression
}

fn is_luau_bytecode(source: &[u8]) -> bool {
    match source.first().copied() {
        None => false,
        Some(version) if version < b'\t' => true,
        Some(version) if version <= BYTECODE_VERSION_MAX => matches!(
            source.get(1),
            Some(BYTECODE_TYPE_VERSION_MIN..=BYTECODE_TYPE_VERSION_MAX)
        ),
        Some(BYTECODE_VERSION_CLASSES) => matches!(
            source.get(1),
            Some(BYTECODE_TYPE_VERSION_MIN..=BYTECODE_TYPE_VERSION_MAX)
        ),
        Some(_) => false,
    }
}