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;
#[doc(hidden)]
pub struct CapturedChunk<F> {
source: &'static str,
make_environment: Cell<Option<F>>,
}
#[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()))
}
}
#[doc(hidden)]
pub struct CaptureEnvironment<'lua> {
lua: LuaRef<'lua>,
environment: Table<'lua>,
names: Table<'lua>,
globals: Table<'lua>,
}
impl<'lua> CaptureEnvironment<'lua> {
#[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()?,
})
}
#[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)
}
#[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)
}
}
}