luau 0.732.0

Safe lifetime-bound Rust embedding API for the Luau runtime
use std::borrow::Cow;
use std::cell::Cell;
use std::io;

use luau_vm::native::{NativeCallContext, NativeCallResult};
use luau_vm::thread::{StackGuard, upvalue_index};
use luau_vm::types::LUA_TNIL;

use super::{AsChunk, ChunkMode};
use crate::error::Error;
use crate::lua::LuaRef;
use crate::table::Table;
use crate::value::IntoLua;

/// A source chunk whose captures are materialized into a private environment
/// when it is passed to [`crate::Lua::load`].
#[doc(hidden)]
pub struct CapturedChunk<F> {
    source: &'static str,
    make_environment: Cell<Option<F>>,
}

/// Creates the chunk representation used by [`crate::chunk!`].
#[doc(hidden)]
pub fn captured_chunk<F>(source: &'static str, make_environment: F) -> CapturedChunk<F> {
    CapturedChunk {
        source,
        make_environment: Cell::new(Some(make_environment)),
    }
}

impl<'lua, F> AsChunk<'lua> for CapturedChunk<F>
where
    F: FnOnce(LuaRef<'lua>) -> Result<Table<'lua>, Error>,
{
    fn environment(&self, lua: &LuaRef<'lua>) -> Result<Option<Table<'lua>>, Error> {
        match self.make_environment.take() {
            Some(make_environment) => make_environment(*lua).map(Some),
            None => Ok(None),
        }
    }

    fn mode(&self) -> Option<ChunkMode> {
        Some(ChunkMode::Text)
    }

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

/// Builds the capture environment used by [`crate::chunk!`].
#[doc(hidden)]
pub struct CaptureEnvironment<'lua> {
    lua: LuaRef<'lua>,
    environment: Table<'lua>,
    names: Table<'lua>,
    globals: Table<'lua>,
}

impl<'lua> CaptureEnvironment<'lua> {
    /// Creates an empty capture environment.
    #[doc(hidden)]
    pub fn new(lua: LuaRef<'lua>) -> Result<Self, Error> {
        Ok(Self {
            lua,
            environment: lua.create_table()?,
            names: lua.create_table()?,
            globals: lua.globals()?,
        })
    }

    /// Adds one captured binding.
    #[doc(hidden)]
    pub fn insert(&self, name: &'static str, value: impl IntoLua<'lua>) -> Result<(), Error> {
        self.names.raw_set(name, true)?;
        self.environment.raw_set(name, value)
    }

    /// Completes the environment and installs its global lookup behavior.
    #[doc(hidden)]
    pub fn finish(self) -> Result<Table<'lua>, Error> {
        let metatable = self.lua.create_table()?;
        let thread = self.lua.current_thread();
        let vm_thread = self.lua.as_vm();

        unsafe {
            let _stack = StackGuard::new(vm_thread);
            metatable.push_to(&thread)?;
            let metatable_index = vm_thread.get_top();
            self.names.push_to(&thread)?;
            self.globals.push_to(&thread)?;
            vm_thread
                .push_native_closure(
                    capture_environment_access,
                    Some("__luau_chunk_environment"),
                    2,
                )
                .map_err(|error| Error::from_thread_exit(vm_thread, error))?;
            vm_thread
                .push_value(-1)
                .map_err(|error| Error::from_thread_exit(vm_thread, error))?;
            vm_thread
                .raw_set_field(metatable_index, "__index")
                .map_err(|error| Error::from_thread_exit(vm_thread, error))?;
            vm_thread
                .raw_set_field(metatable_index, "__newindex")
                .map_err(|error| Error::from_thread_exit(vm_thread, error))?;
        }

        self.environment.set_metatable(Some(&metatable))?;
        Ok(self.environment)
    }
}

fn capture_environment_access(context: NativeCallContext<'_>) -> NativeCallResult {
    unsafe {
        let thread = context.raw_thread();
        let arg_count = context.arg_count();
        if !matches!(arg_count, 2 | 3) {
            return luau_vm::error!(thread, "invalid captured chunk environment access")
                .map_err(Into::into);
        }
        if thread.is_table(upvalue_index(1)) == 0 || thread.is_table(upvalue_index(2)) == 0 {
            return luau_vm::error!(thread, "invalid captured chunk environment state")
                .map_err(Into::into);
        }

        thread.push_value(2)?;
        let captured = thread.raw_get(upvalue_index(1)) != LUA_TNIL;
        thread.pop(1);

        if arg_count == 2 {
            if captured {
                thread.push_nil()?;
            } else {
                thread.push_value(2)?;
                thread.get_table(upvalue_index(2))?;
            }
            Ok(1)
        } else {
            if captured && thread.is_table(1) == 0 {
                return context.arg(1).type_error("table").map_err(Into::into);
            }
            thread.push_value(2)?;
            thread.push_value(3)?;
            if captured {
                thread.raw_set(1)?;
            } else {
                thread.set_table(upvalue_index(2))?;
            }
            Ok(0)
        }
    }
}