luau 0.732.0

use core::fmt;
use core::ops::{Deref, DerefMut};
use std::cell::{Ref, RefCell, RefMut};

use luau_common::ByteSlice;
use luau_vm::internal::api::RawStackAccess;
use luau_vm::internal::userdata::{TypedUserdata, TypedUserdataAccess, Userdata as VmUserdata};
use luau_vm::thread::{LUA_MULTRET, StackGuard, Thread as VmThread};
use luau_vm::types::LUA_TUSERDATA;

use super::registration::UserdataProxy;
use super::{MetaMethod, Userdata};
use crate::error::Error;
use crate::function::Function;
use crate::object::{self, ObjectLike};
use crate::string::LuaString;
use crate::table::{Table, TablePairs};
use crate::thread::Thread;
use crate::value::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti, Value, ValueRef};

/// A handle to a Luau userdata value.
///
/// Registered Rust payloads can be checked and borrowed through this handle.
pub struct AnyUserdata<'lua> {
    reference: ValueRef<'lua>,
    userdata: VmUserdata,
}

/// Shared borrow of a registered Rust userdata payload.
pub struct UserdataRef<'lua, T> {
    value: Ref<'lua, T>,
}

/// Exclusive borrow of a registered Rust userdata payload.
pub struct UserdataRefMut<'lua, T> {
    value: RefMut<'lua, T>,
}

/// Restricted access to the metatable shared by a registered userdata type.
pub struct UserdataMetatable<'lua> {
    table: Table<'lua>,
}

/// Iterator over accessible entries in a registered userdata metatable.
pub struct UserdataMetatablePairs<'table, 'lua, V> {
    pairs: TablePairs<'table, 'lua, LuaString<'lua>, V>,
}

impl<'lua> AnyUserdata<'lua> {
    pub(crate) fn from_stack(thread: &Thread<'lua>, index: i32) -> Result<Self, Error> {
        if unsafe { thread.as_vm().type_of(index) } != LUA_TUSERDATA {
            return Err(Error::from_lua_conversion(
                thread.stack_type_name(index).as_str(),
                "userdata",
                None,
            ));
        }
        let userdata = unsafe {
            thread
                .as_vm()
                .to_object(index)
                .expect("userdata stack slot should contain userdata")
                .userdata_value()
        };
        Ok(Self {
            reference: ValueRef::from_stack(thread, index)
                .map_err(|exit| Error::from_thread_exit(thread, exit))?,
            userdata,
        })
    }

    /// Creates another rooted handle to this userdata.
    pub fn try_clone(&self) -> Result<Self, Error> {
        Ok(Self {
            reference: self.reference.try_clone()?,
            userdata: self.userdata,
        })
    }

    /// Compares two userdata values using Luau's equality semantics.
    ///
    /// This may invoke the `__eq` metamethod.
    pub fn equals(&self, other: &Self) -> Result<bool, Error> {
        if self == other {
            return Ok(true);
        }

        self.ensure_usable()?;
        other.ensure_usable()?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            other.push_to(&thread)?;
            vm_thread
                .equal(-2, -1)
                .map(|equal| equal != 0)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Returns whether this userdata contains a Rust value of type `T`.
    pub fn is<T: 'static>(&self) -> bool {
        self.typed().is_some_and(|userdata| userdata.is::<T>())
    }

    /// Returns whether this is the type-level proxy for `T`.
    pub fn is_proxy<T: 'static>(&self) -> bool {
        self.is::<UserdataProxy<T>>()
    }

    /// Immutably borrows the Rust payload as `T`.
    pub fn borrow<T: 'static>(&self) -> Result<UserdataRef<'_, T>, Error> {
        let cell = self.cell::<T>()?;
        let value = cell.try_borrow().map_err(|_| Error::UserdataBorrowError)?;
        let value =
            Ref::filter_map(value, Option::as_ref).map_err(|_| Error::UserdataDestructed)?;
        Ok(UserdataRef { value })
    }

    /// Mutably borrows the Rust payload as `T`.
    pub fn borrow_mut<T: 'static>(&self) -> Result<UserdataRefMut<'_, T>, Error> {
        let cell = self.cell::<T>()?;
        let value = cell
            .try_borrow_mut()
            .map_err(|_| Error::UserdataBorrowMutError)?;
        let value =
            RefMut::filter_map(value, Option::as_mut).map_err(|_| Error::UserdataDestructed)?;
        self.reference.runtime().invalidate_managed_safe_env();
        Ok(UserdataRefMut { value })
    }

    /// Moves the registered Rust payload out and invalidates the userdata.
    pub fn take<T: 'static>(&self) -> Result<T, Error> {
        let Some(userdata) = self.typed() else {
            return Err(Error::UserdataTypeMismatch);
        };
        let value = unsafe {
            self.reference
                .reference_thread()
                .take_typed_userdata::<T>(userdata)
                .map_err(Error::from_typed_userdata)
        }?;
        self.reference.runtime().invalidate_managed_safe_env();
        Ok(value)
    }

    /// Drops the registered Rust payload and invalidates the userdata.
    pub fn destroy(&self) -> Result<(), Error> {
        let Some(userdata) = self.typed() else {
            return Err(Error::UserdataTypeMismatch);
        };
        self.reference.runtime().invalidate_managed_safe_env();
        unsafe {
            self.reference
                .reference_thread()
                .destroy_typed_userdata(userdata)
                .map_err(Error::from_typed_userdata)
        }?;
        Ok(())
    }

    /// Associates a Lua value with this registered Rust userdata.
    ///
    /// The value is traced from the userdata and is released when replaced or
    /// when the userdata is collected. Taking or destroying the Rust payload
    /// leaves the associated value attached to the userdata. Store a table
    /// here when an instance needs more than one associated value.
    ///
    /// Associated values remain replaceable after the Rust payload has been
    /// taken or destroyed.
    pub fn set_user_value(&self, value: impl IntoLua<'lua>) -> Result<(), Error> {
        let Some(userdata) = self.typed() else {
            return Err(Error::UserdataTypeMismatch);
        };

        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            value.push_into_stack(&thread)?;
            let value = vm_thread
                .to_object(-1)
                .expect("pushed userdata value should occupy a stack slot");
            self.reference.runtime().invalidate_managed_safe_env();
            vm_thread.set_typed_userdata_value(userdata, value);
            Ok(())
        }
    }

    /// Returns the Lua value associated with this registered Rust userdata.
    ///
    /// The associated value remains readable after the Rust payload has been
    /// taken or destroyed.
    pub fn user_value<V>(&self) -> Result<V, Error>
    where
        V: FromLua<'lua>,
    {
        let Some(userdata) = self.typed() else {
            return Err(Error::UserdataTypeMismatch);
        };

        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            let value = vm_thread.typed_userdata_value(userdata);
            vm_thread
                .push_value_internal(value)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            V::from_stack(&thread, -1)
        }
    }

    /// Returns restricted access to this userdata's metatable.
    pub fn metatable(&self) -> Result<UserdataMetatable<'lua>, Error> {
        let Some(userdata) = self.typed() else {
            return Err(Error::UserdataTypeMismatch);
        };
        if userdata.is_destructed() {
            return Err(Error::UserdataDestructed);
        }

        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            let metatable = self
                .userdata
                .metatable()
                .expect("registered userdata must have a metatable");
            vm_thread
                .push_table(metatable)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            let table = Table::from_stack(&thread, -1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            Ok(UserdataMetatable { table })
        }
    }

    /// Returns this userdata's `__type` name, or `"userdata"`.
    pub fn type_name(&self) -> Result<LuaString<'lua>, Error> {
        self.ensure_usable()?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            self.push_to(&thread)?;
            if vm_thread
                .get_metafield(-1, MetaMethod::Type.name())
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?
                == 0
            {
                vm_thread
                    .push_string("userdata")
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            }
            if vm_thread.is_string(-1) == 0 {
                vm_thread.pop(1);
                vm_thread
                    .push_string("userdata")
                    .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            }
            LuaString::from_stack(&thread, -1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))
        }
    }

    /// Returns a pointer that uniquely identifies this userdata.
    pub fn to_pointer(&self) -> *const () {
        self.pointer()
    }

    pub(crate) fn push_to(&self, target: impl AsRef<VmThread>) -> Result<(), Error> {
        self.reference.push_to(target)
    }

    pub(crate) fn thread(&self) -> Thread<'lua> {
        self.reference.thread()
    }

    pub(crate) fn pointer(&self) -> *const () {
        self.reference.pointer()
    }

    pub(crate) fn fmt_pretty(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        if let Some(value) = self.debug_string() {
            return formatter.write_str(&value);
        }

        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            self.push_to(&thread).map_err(|_| fmt::Error)?;
            let type_name = vm_thread.lua_type_name(-1);
            write!(
                formatter,
                "{}: {:p}",
                type_name.as_bytes().to_str_lossy(),
                self.pointer()
            )
        }
    }

    fn debug_string(&self) -> Option<String> {
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            self.push_to(&thread).ok()?;

            let mut called = vm_thread
                .call_meta(-1, MetaMethod::ToDebugString.name())
                .ok()?;
            if called == 0 {
                called = vm_thread.call_meta(-1, MetaMethod::ToString.name()).ok()?;
            }
            if called == 0 {
                return None;
            }

            vm_thread
                .to_string(-1)
                .ok()
                .flatten()
                .map(ToString::to_string)
        }
    }

    fn cell<T: 'static>(&self) -> Result<&RefCell<Option<T>>, Error> {
        let Some(userdata) = self.typed() else {
            return Err(Error::UserdataTypeMismatch);
        };
        if userdata.is_destructed() {
            return Err(Error::UserdataDestructed);
        }
        let Some(cell) = (unsafe { userdata.cell_ptr::<T>() }) else {
            return Err(Error::UserdataTypeMismatch);
        };

        Ok(unsafe { &*cell })
    }

    fn typed(&self) -> Option<TypedUserdata> {
        unsafe {
            self.reference
                .reference_thread()
                .typed_userdata(self.userdata)
        }
    }

    fn ensure_usable(&self) -> Result<(), Error> {
        if self
            .typed()
            .is_some_and(|userdata| userdata.is_destructed())
        {
            return Err(Error::UserdataDestructed);
        }
        Ok(())
    }
}

impl<'lua> UserdataMetatable<'lua> {
    /// Creates another rooted handle to this metatable.
    pub fn try_clone(&self) -> Result<Self, Error> {
        Ok(Self {
            table: self.table.try_clone()?,
        })
    }

    /// Gets an accessible metatable entry.
    pub fn get<V>(&self, key: impl AsRef<[u8]>) -> Result<V, Error>
    where
        V: FromLua<'lua>,
    {
        MetaMethod::validate(key.as_ref())?;
        self.table.raw_get(key.as_ref())
    }

    /// Sets an accessible metatable entry.
    ///
    /// The protected `__gc`, `__metatable`, `__index`, and `__newindex`
    /// entries cannot be replaced through this handle.
    pub fn set(&self, key: impl AsRef<[u8]>, value: impl IntoLua<'lua>) -> Result<(), Error> {
        let key = key.as_ref();
        MetaMethod::validate(key)?;
        if key == MetaMethod::Index.name().as_bytes()
            || key == MetaMethod::NewIndex.name().as_bytes()
        {
            return Err(Error::MetaMethodRestricted(
                String::from_utf8_lossy(key).into_owned(),
            ));
        }
        self.table.raw_set(key, value)
    }

    /// Returns whether an accessible metatable entry is present.
    pub fn contains(&self, key: impl AsRef<[u8]>) -> Result<bool, Error> {
        MetaMethod::validate(key.as_ref())?;
        self.table
            .raw_get::<Value<'lua>>(key.as_ref())
            .map(|value| !value.is_nil())
    }

    /// Iterates over accessible metatable entries.
    pub fn pairs<V>(&self) -> UserdataMetatablePairs<'_, 'lua, V>
    where
        V: FromLua<'lua>,
    {
        UserdataMetatablePairs {
            pairs: self.table.pairs(),
        }
    }
}

impl<'lua, V> Iterator for UserdataMetatablePairs<'_, 'lua, V>
where
    V: FromLua<'lua>,
{
    type Item = Result<(LuaString<'lua>, V), Error>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.pairs.next()? {
                Ok((key, value)) if MetaMethod::validate(key.as_bytes()).is_ok() => {
                    return Some(Ok((key, value)));
                }
                Ok(_) => {}
                Err(error) => return Some(Err(error)),
            }
        }
    }
}

impl<'lua, 'userdata> IntoLua<'lua> for AnyUserdata<'userdata>
where
    'userdata: 'lua,
{
    fn into_lua(self, _: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        Ok(Value::Userdata(self))
    }

    unsafe fn push_into_stack(self, thread: &Thread<'lua>) -> Result<(), Error> {
        self.push_to(thread)
    }
}

impl<'lua, 'userdata> IntoLua<'lua> for &AnyUserdata<'userdata>
where
    'userdata: 'lua,
{
    fn into_lua(self, _: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        Ok(Value::Userdata(self.try_clone()?))
    }

    unsafe fn push_into_stack(self, thread: &Thread<'lua>) -> Result<(), Error> {
        self.push_to(thread)
    }
}

impl<'lua> FromLua<'lua> for AnyUserdata<'lua> {
    fn from_lua(value: Value<'lua>, _: crate::LuaRef<'lua>) -> Result<Self, Error> {
        match value {
            Value::Userdata(userdata) => Ok(userdata),
            value => Err(Error::from_lua_conversion(
                value.type_name(),
                "userdata",
                None,
            )),
        }
    }
}

impl<'lua, T> IntoLua<'lua> for T
where
    T: Userdata + 'static,
{
    fn into_lua(self, thread: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        thread.create_userdata(self).map(Value::Userdata)
    }

    unsafe fn push_into_stack(self, thread: &Thread<'lua>) -> Result<(), Error> {
        thread.push_userdata(self)
    }
}

impl PartialEq for AnyUserdata<'_> {
    fn eq(&self, other: &Self) -> bool {
        self.reference == other.reference
    }
}

impl Eq for AnyUserdata<'_> {}

impl object::private::Sealed for AnyUserdata<'_> {}

impl<'lua> ObjectLike<'lua> for AnyUserdata<'lua> {
    fn get<V>(&self, key: impl IntoLua<'lua>) -> Result<V, Error>
    where
        V: FromLua<'lua>,
    {
        self.ensure_usable()?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            let userdata_index = vm_thread.get_top();
            key.push_into_stack(&thread)?;
            vm_thread
                .get_table(userdata_index)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            V::from_stack(&thread, -1)
        }
    }

    fn set(&self, key: impl IntoLua<'lua>, value: impl IntoLua<'lua>) -> Result<(), Error> {
        self.ensure_usable()?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            let userdata_index = vm_thread.get_top();
            key.push_into_stack(&thread)?;
            value.push_into_stack(&thread)?;
            self.reference.runtime().invalidate_managed_safe_env();
            vm_thread
                .set_table(userdata_index)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
        }
        Ok(())
    }

    fn call<R>(&self, args: impl IntoLuaMulti<'lua>) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        self.ensure_usable()?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let stack = StackGuard::new(vm_thread);

            self.push_to(&thread)?;
            let arg_count = i32::try_from(args.push_into_stack_multi(&thread)?)
                .map_err(|_| Error::StackError)?;
            vm_thread
                .protected_call(arg_count, LUA_MULTRET, 0)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            let result_count = vm_thread.get_top() - stack.top();
            R::from_stack_multi(&thread, stack.top(), result_count)
        }
    }

    fn call_method<R>(&self, name: &str, args: impl IntoLuaMulti<'lua>) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        let function = self.get::<Function<'lua>>(name)?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let stack = StackGuard::new(vm_thread);

            function.push_to(&thread)?;
            self.push_to(&thread)?;
            let arg_count = args
                .push_into_stack_multi(&thread)?
                .checked_add(1)
                .ok_or(Error::StackError)?;
            let arg_count = i32::try_from(arg_count).map_err(|_| Error::StackError)?;
            vm_thread
                .protected_call(arg_count, LUA_MULTRET, 0)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            let result_count = vm_thread.get_top() - stack.top();
            R::from_stack_multi(&thread, stack.top(), result_count)
        }
    }

    fn call_function<R>(&self, name: &str, args: impl IntoLuaMulti<'lua>) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        self.get::<Function<'lua>>(name)?.call(args)
    }

    fn to_string(&self) -> Result<String, Error> {
        self.ensure_usable()?;
        unsafe {
            let thread = self.thread();
            let vm_thread = thread.as_vm();
            let _stack = StackGuard::new(vm_thread);
            self.push_to(&thread)?;
            let bytes = vm_thread
                .lua_to_string(-1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            bytes.to_str().map(str::to_owned).map_err(|error| {
                let message = error.to_string();
                Error::from_lua_conversion("string", "String", Some(message.as_str()))
            })
        }
    }

    fn to_value(&self) -> Result<Value<'lua>, Error> {
        self.try_clone().map(Value::Userdata)
    }
}

impl fmt::Debug for AnyUserdata<'_> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        if formatter.alternate() {
            return self.fmt_pretty(formatter);
        }

        formatter
            .debug_tuple("AnyUserdata")
            .field(&self.reference)
            .finish()
    }
}

impl<T> Deref for UserdataRef<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<T> Deref for UserdataRefMut<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<T> DerefMut for UserdataRefMut<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.value
    }
}