luau 0.732.0

use core::any::TypeId;
use core::cell::RefCell;
use core::marker::PhantomData;
use std::collections::HashMap;

use luau_vm::internal::userdata::{
    TypedUserdataAccess, UserdataTypeRegistration, UserdataTypeRegistryAccess,
};
use luau_vm::thread::{StackGuard, Thread as VmThread};

use super::registry::UserdataDefinition;
use super::traits::short_type_name;
use super::{AnyUserdata, Userdata, UserdataRegistry};
use crate::error::Error;
use crate::lua::{Lua, LuaRef};
use crate::thread::Thread;

pub(super) struct UserdataProxy<T>(PhantomData<fn() -> T>);

#[derive(Default)]
pub(crate) struct PendingUserdataRegistrations {
    definitions: RefCell<HashMap<TypeId, UserdataDefinition>>,
}

impl<T: Userdata + 'static> Userdata for UserdataProxy<T> {
    fn register(registry: &mut UserdataRegistry<'_, Self>) {
        registry.register_proxy_target::<T>();
    }
}

impl Lua {
    /// Creates userdata from a type whose fields and methods come from its
    /// [`Userdata`] implementation.
    pub fn create_userdata<T: Userdata + 'static>(
        &self,
        value: T,
    ) -> Result<AnyUserdata<'_>, Error> {
        self.lua_ref().create_userdata(value)
    }

    /// Creates userdata for any `'static` Rust value.
    ///
    /// Use [`Lua::register_userdata_type`] to configure fields and methods.
    /// Otherwise the type receives an empty registration.
    pub fn create_any_userdata<T: 'static>(&self, value: T) -> Result<AnyUserdata<'_>, Error> {
        self.lua_ref().create_any_userdata(value)
    }

    /// Registers fields and methods for arbitrary userdata of type `T`.
    ///
    /// Re-registering a type replaces the registration used by subsequently
    /// created userdata. Existing userdata retain their original metatable.
    pub fn register_userdata_type<T: 'static>(
        &self,
        register: impl FnOnce(&mut UserdataRegistry<'_, T>),
    ) -> Result<(), Error> {
        self.lua_ref().register_userdata_type(register)
    }

    /// Creates the type-level proxy for `T`.
    ///
    /// The proxy exposes fields and functions registered without an instance
    /// receiver, allowing it to serve as a constructor or namespace in Luau.
    pub fn create_proxy<T: Userdata + 'static>(&self) -> Result<AnyUserdata<'_>, Error> {
        self.lua_ref().create_proxy::<T>()
    }
}

impl<'lua> LuaRef<'lua> {
    /// Creates userdata from a type whose fields and methods come from its
    /// [`Userdata`] implementation.
    pub fn create_userdata<T: Userdata + 'static>(
        &self,
        value: T,
    ) -> Result<AnyUserdata<'lua>, Error> {
        self.current_thread().create_userdata(value)
    }

    /// Creates userdata for any `'static` Rust value.
    ///
    /// An empty shared registration is created if the type was not explicitly
    /// registered first.
    pub fn create_any_userdata<T: 'static>(&self, value: T) -> Result<AnyUserdata<'lua>, Error> {
        self.current_thread().create_any_userdata(value)
    }

    /// Registers fields and methods for arbitrary userdata of type `T`.
    ///
    /// Re-registering a type replaces the registration used by subsequently
    /// created userdata. Existing userdata retain their original metatable.
    pub fn register_userdata_type<T: 'static>(
        &self,
        register: impl FnOnce(&mut UserdataRegistry<'_, T>),
    ) -> Result<(), Error> {
        self.current_thread().register_userdata_type(register)
    }

    /// Creates the type-level proxy for `T`.
    pub fn create_proxy<T: Userdata + 'static>(&self) -> Result<AnyUserdata<'lua>, Error> {
        self.current_thread().create_proxy::<T>()
    }
}

impl<'lua> Thread<'lua> {
    pub(crate) fn create_userdata<T: Userdata + 'static>(
        &self,
        value: T,
    ) -> Result<AnyUserdata<'lua>, Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            self.push_userdata_with(value, T::register)?;
            AnyUserdata::from_stack(self, -1)
        }
    }

    pub(crate) fn create_any_userdata<T: 'static>(
        &self,
        value: T,
    ) -> Result<AnyUserdata<'lua>, Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            self.push_userdata_with(value, |_| {})?;
            AnyUserdata::from_stack(self, -1)
        }
    }

    pub(crate) fn register_userdata_type<T: 'static>(
        &self,
        register: impl FnOnce(&mut UserdataRegistry<'_, T>),
    ) -> Result<(), Error> {
        let definition = capture_userdata_definition(self, register)?;
        self.runtime()
            .userdata_registrations()
            .replace::<T>(self.reference_thread(), definition);
        unsafe { self.reference_thread().retire_userdata_type::<T>() };
        Ok(())
    }

    pub(crate) fn create_proxy<T: Userdata + 'static>(&self) -> Result<AnyUserdata<'lua>, Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            let registration = self.userdata_registration::<UserdataProxy<T>>(
                short_type_name::<T>,
                UserdataProxy::<T>::register,
            )?;
            thread
                .push_typed_userdata(UserdataProxy::<T>(PhantomData), &registration)
                .map_err(|error| Error::from_thread_exit(thread, error))?;
            AnyUserdata::from_stack(self, -1)
        }
    }

    pub(crate) fn push_userdata<T: Userdata + 'static>(&self, value: T) -> Result<(), Error> {
        self.push_userdata_with(value, T::register)
    }

    fn push_userdata_with<T: 'static>(
        &self,
        value: T,
        register: impl FnOnce(&mut UserdataRegistry<'_, T>),
    ) -> Result<(), Error> {
        let thread = self.as_vm();
        let registration = self.userdata_registration::<T>(short_type_name::<T>, register)?;
        unsafe {
            thread
                .push_typed_userdata(value, &registration)
                .map_err(|error| Error::from_thread_exit(thread, error))
        }
    }

    fn userdata_registration<T: 'static>(
        &self,
        name: impl FnOnce() -> String,
        register: impl FnOnce(&mut UserdataRegistry<'_, T>),
    ) -> Result<UserdataTypeRegistration, Error> {
        if let Some(definition) = self.runtime().userdata_registrations().take::<T>() {
            return definition.materialize::<T>(self, name());
        }

        let vm_thread = self.reference_thread();
        unsafe {
            if let Some(registration) = vm_thread.userdata_type::<T>() {
                return Ok(registration);
            }
        }

        capture_userdata_definition(self, register)?.materialize::<T>(self, name())
    }
}

impl PendingUserdataRegistrations {
    fn replace<T: 'static>(&self, thread: &VmThread, definition: UserdataDefinition) {
        if let Some(previous) = self
            .definitions
            .borrow_mut()
            .insert(TypeId::of::<T>(), definition)
        {
            previous.release(thread);
        }
    }

    fn take<T: 'static>(&self) -> Option<UserdataDefinition> {
        self.definitions.borrow_mut().remove(&TypeId::of::<T>())
    }

    pub(crate) fn close(&self, thread: &VmThread) {
        for (_, definition) in self.definitions.take() {
            definition.release(thread);
        }
    }
}

fn capture_userdata_definition<T>(
    thread: &Thread<'_>,
    register: impl FnOnce(&mut UserdataRegistry<'_, T>),
) -> Result<UserdataDefinition, Error> {
    let mut registry = UserdataRegistry::<T>::new(thread.reference_thread(), thread.runtime());
    register(&mut registry);
    registry.finish()
}