luau 0.732.0

use std::fmt;
use std::hash::{Hash, Hasher};
use std::mem;
use std::sync::{Arc, Mutex};

use luau_vm::Thread as VmThread;
use luau_vm::thread::{LUA_REFNIL, LUA_REGISTRY_INDEX, StackGuard};
use luau_vm::types::LUA_TNIL;

use super::{Lua, LuaRef};
use crate::error::Error;
use crate::thread::Thread;
use crate::value::{FromLua, IntoLua, Value};

#[derive(Clone, Default)]
pub(crate) struct RegistryState {
    unref_list: Arc<Mutex<Option<Vec<i32>>>>,
}

/// A key for a value stored in the Luau registry.
///
/// Dropping a key defers removal of its registry value until
/// [`Lua::expire_registry_values`] is called. A key can instead be removed
/// immediately with [`Lua::remove_registry_value`].
///
/// # Reference cycles
///
/// Storing a `RegistryKey` inside a [`crate::Userdata`] value or Rust callback
/// can create a cycle that Luau's garbage collector cannot resolve. For
/// example, a userdata that owns a key referring back to itself remains rooted
/// by that registry entry, so its Rust payload never reaches `Drop`. Remove
/// such keys explicitly. Use [`crate::AnyUserdata::set_user_value`] to
/// associate Luau values with userdata without creating a permanent root.
pub struct RegistryKey {
    reference: i32,
    unref_list: Arc<Mutex<Option<Vec<i32>>>>,
}

impl Lua {
    /// Stores a value in the registry under a generated key.
    pub fn create_registry_value<'lua>(
        &'lua self,
        value: impl IntoLua<'lua>,
    ) -> Result<RegistryKey, Error> {
        self.lua_ref().create_registry_value(value)
    }

    /// Stores a value in the registry under a byte-string key.
    pub fn set_named_registry_value<'lua>(
        &'lua self,
        key: impl AsRef<[u8]>,
        value: impl IntoLua<'lua>,
    ) -> Result<(), Error> {
        self.lua_ref().set_named_registry_value(key, value)
    }

    /// Retrieves a named registry value.
    pub fn named_registry_value<'lua, T>(&'lua self, key: impl AsRef<[u8]>) -> Result<T, Error>
    where
        T: FromLua<'lua>,
    {
        self.lua_ref().named_registry_value(key)
    }

    /// Removes a named registry value.
    pub fn unset_named_registry_value(&self, key: impl AsRef<[u8]>) -> Result<(), Error> {
        self.lua_ref().unset_named_registry_value(key)
    }

    /// Retrieves the value associated with a generated registry key.
    pub fn registry_value<'lua, T>(&'lua self, key: &RegistryKey) -> Result<T, Error>
    where
        T: FromLua<'lua>,
    {
        self.lua_ref().registry_value(key)
    }

    /// Removes a generated registry value immediately.
    pub fn remove_registry_value(&self, key: RegistryKey) -> Result<(), Error> {
        if !self.owns_registry_value(&key) {
            return Err(Error::mismatched_registry_key());
        }

        let reference = key.take();
        unsafe {
            self.state.main_thread().unref_value(reference);
        }
        Ok(())
    }

    /// Replaces the value associated with a generated registry key.
    pub fn replace_registry_value<'lua>(
        &'lua self,
        key: &mut RegistryKey,
        value: impl IntoLua<'lua>,
    ) -> Result<(), Error> {
        if !self.owns_registry_value(key) {
            return Err(Error::mismatched_registry_key());
        }

        self.lua_ref().replace_registry_value(key, value)
    }

    /// Returns whether this state created the registry key.
    pub fn owns_registry_value(&self, key: &RegistryKey) -> bool {
        self.lua_ref().owns_registry_value(key)
    }

    /// Removes registry values whose keys have been dropped.
    pub fn expire_registry_values(&self) {
        self.lua_ref().expire_registry_values();
    }
}

impl<'lua> LuaRef<'lua> {
    /// Stores a value in the registry under a generated key.
    pub fn create_registry_value(&self, value: impl IntoLua<'lua>) -> Result<RegistryKey, Error> {
        self.current_thread().create_registry_value(value)
    }

    /// Stores a value in the registry under a byte-string key.
    pub fn set_named_registry_value(
        &self,
        key: impl AsRef<[u8]>,
        value: impl IntoLua<'lua>,
    ) -> Result<(), Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            value.push_into_stack(&self.current_thread())?;
            thread
                .raw_set_field(LUA_REGISTRY_INDEX, key)
                .map_err(|exit| Error::from_thread_exit(thread, exit))
        }
    }

    /// Retrieves a named registry value.
    pub fn named_registry_value<T>(&self, key: impl AsRef<[u8]>) -> Result<T, Error>
    where
        T: FromLua<'lua>,
    {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            let safe_thread = self.current_thread();
            thread
                .raw_get_field(LUA_REGISTRY_INDEX, key)
                .map_err(|exit| Error::from_thread_exit(thread, exit))?;
            T::from_stack(&safe_thread, -1)
        }
    }

    /// Removes a named registry value.
    pub fn unset_named_registry_value(&self, key: impl AsRef<[u8]>) -> Result<(), Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            thread
                .push_nil()
                .map_err(|exit| Error::from_thread_exit(thread, exit))?;
            thread
                .raw_set_field(LUA_REGISTRY_INDEX, key)
                .map_err(|exit| Error::from_thread_exit(thread, exit))
        }
    }

    /// Retrieves the value associated with a generated registry key.
    pub fn registry_value<T>(&self, key: &RegistryKey) -> Result<T, Error>
    where
        T: FromLua<'lua>,
    {
        self.current_thread().registry_value(key)
    }

    /// Replaces the value associated with a generated registry key.
    pub fn replace_registry_value(
        &self,
        key: &mut RegistryKey,
        value: impl IntoLua<'lua>,
    ) -> Result<(), Error> {
        if !self.owns_registry_value(key) {
            return Err(Error::mismatched_registry_key());
        }

        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            let safe_thread = self.current_thread();

            value.push_into_stack(&safe_thread)?;
            let is_nil = thread.type_of(-1) == LUA_TNIL;
            match (key.id(), is_nil) {
                (LUA_REFNIL, true) => {}
                (reference, true) => {
                    thread.unref_value(reference);
                    key.set_id(LUA_REFNIL);
                }
                (LUA_REFNIL, false) => {
                    let replacement = safe_thread.create_registry_key_from_top()?;
                    key.set_id(replacement.take());
                }
                (reference, false) => {
                    thread
                        .raw_seti(LUA_REGISTRY_INDEX, reference)
                        .map_err(|exit| Error::from_thread_exit(thread, exit))?;
                }
            }
        }
        Ok(())
    }

    /// Returns whether this state created the registry key.
    pub fn owns_registry_value(&self, key: &RegistryKey) -> bool {
        self.runtime().registry().owns(key)
    }

    /// Removes registry values whose keys have been dropped.
    pub fn expire_registry_values(&self) {
        self.runtime().registry().expire(self.as_vm());
    }
}

impl<'lua> Thread<'lua> {
    pub(crate) fn create_registry_value(
        &self,
        value: impl IntoLua<'lua>,
    ) -> Result<RegistryKey, Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            value.push_into_stack(self)?;
            self.create_registry_key_from_top()
        }
    }

    /// # Safety
    ///
    /// The current stack top must contain the value to store.
    pub(crate) unsafe fn create_registry_key_from_top(&self) -> Result<RegistryKey, Error> {
        unsafe {
            let thread = self.as_vm();
            let reference = if thread.type_of(-1) == LUA_TNIL {
                LUA_REFNIL
            } else if let Some(reference) = self.registry().take_dropped_reference() {
                if let Err(exit) = thread.raw_seti(LUA_REGISTRY_INDEX, reference) {
                    self.registry().restore_dropped_reference(reference);
                    return Err(Error::from_thread_exit(thread, exit));
                }
                reference
            } else {
                thread
                    .ref_value(-1)
                    .map_err(|exit| Error::from_thread_exit(thread, exit))?
            };
            Ok(self.registry().create_key(reference))
        }
    }

    pub(crate) fn registry_value<T>(&self, key: &RegistryKey) -> Result<T, Error>
    where
        T: FromLua<'lua>,
    {
        if !self.registry().owns(key) {
            return Err(Error::mismatched_registry_key());
        }

        if key.id() == LUA_REFNIL {
            return T::from_lua(Value::Nil, self.lua_ref());
        }

        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            thread
                .get_ref(key.id())
                .map_err(|exit| Error::from_thread_exit(thread, exit))?;
            T::from_stack(self, -1)
        }
    }

    pub(crate) fn push_registry_value(&self, key: &RegistryKey) -> Result<(), Error> {
        if !self.registry().owns(key) {
            return Err(Error::mismatched_registry_key());
        }

        unsafe {
            if key.id() == LUA_REFNIL {
                self.as_vm()
                    .push_nil()
                    .map_err(|exit| Error::from_thread_exit(self, exit))?;
            } else {
                self.as_vm()
                    .get_ref(key.id())
                    .map_err(|exit| Error::from_thread_exit(self, exit))?;
            }
        }
        Ok(())
    }
}

impl RegistryState {
    pub(crate) fn create_key(&self, reference: i32) -> RegistryKey {
        RegistryKey {
            reference,
            unref_list: Arc::clone(&self.unref_list),
        }
    }

    pub(crate) fn owns(&self, key: &RegistryKey) -> bool {
        Arc::ptr_eq(&self.unref_list, &key.unref_list)
    }

    pub(crate) fn take_dropped_reference(&self) -> Option<i32> {
        lock_unref_list(&self.unref_list).as_mut()?.pop()
    }

    pub(crate) fn restore_dropped_reference(&self, reference: i32) {
        if reference <= LUA_REFNIL {
            return;
        }

        if let Some(unref_list) = lock_unref_list(&self.unref_list).as_mut() {
            unref_list.push(reference);
        }
    }

    pub(crate) fn expire(&self, thread: &VmThread) {
        let references = {
            let mut unref_list = lock_unref_list(&self.unref_list);
            let Some(unref_list) = unref_list.as_mut() else {
                return;
            };
            mem::take(unref_list)
        };

        for reference in references {
            unsafe {
                thread.unref_value(reference);
            }
        }
    }

    pub(crate) fn close(&self, thread: &VmThread) {
        let references = {
            let mut unref_list = lock_unref_list(&self.unref_list);
            unref_list.take().unwrap_or_default()
        };

        for reference in references {
            unsafe {
                thread.unref_value(reference);
            }
        }
    }
}

impl RegistryKey {
    /// Returns the underlying registry reference.
    pub fn id(&self) -> i32 {
        self.reference
    }

    pub(crate) fn set_id(&mut self, reference: i32) {
        self.reference = reference;
    }

    pub(crate) fn take(mut self) -> i32 {
        let reference = self.reference;
        self.reference = LUA_REFNIL;
        reference
    }
}

impl Drop for RegistryKey {
    fn drop(&mut self) {
        if self.reference <= LUA_REFNIL {
            return;
        }

        let mut unref_list = lock_unref_list(&self.unref_list);
        if let Some(unref_list) = unref_list.as_mut() {
            unref_list.push(self.reference);
        }
    }
}

impl fmt::Debug for RegistryKey {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "RegistryKey({})", self.id())
    }
}

impl Hash for RegistryKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.id().hash(state);
    }
}

impl PartialEq for RegistryKey {
    fn eq(&self, other: &Self) -> bool {
        self.id() == other.id() && Arc::ptr_eq(&self.unref_list, &other.unref_list)
    }
}

impl Eq for RegistryKey {}

impl<'lua> IntoLua<'lua> for RegistryKey {
    fn into_lua(self, thread: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        (&self).into_lua(thread)
    }

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

impl<'lua> IntoLua<'lua> for &RegistryKey {
    fn into_lua(self, thread: crate::LuaRef<'lua>) -> Result<Value<'lua>, Error> {
        thread.current_thread().registry_value(self)
    }

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

impl<'lua> FromLua<'lua> for RegistryKey {
    fn from_lua(value: Value<'lua>, thread: crate::LuaRef<'lua>) -> Result<Self, Error> {
        thread.create_registry_value(value)
    }
}

fn lock_unref_list(
    unref_list: &Mutex<Option<Vec<i32>>>,
) -> std::sync::MutexGuard<'_, Option<Vec<i32>>> {
    unref_list
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn assert_send_sync<T: Send + Sync>() {}

    #[test]
    fn registry_key_is_send_sync() {
        assert_send_sync::<RegistryKey>();
    }
}