luau-vm 0.732.0

use core::any::TypeId;
use core::cell::RefCell;
use core::ptr::{self, NonNull};
use core::slice;

use crate::Table;
use crate::VmErrorResult;
use crate::gc::{GcBarrier, GcObject};
use crate::handle::RawHandle;
use crate::handle::sealed::Sealed;
use crate::memory::MemoryRuntime;
use crate::native::{NativeCallContext, NativeCallResult};
use crate::state::GlobalState;
use crate::table::RawLuaTable;
use crate::thread::stack::RawStackAccess;
use crate::thread::{StackGuard, Thread};
use crate::types::{LUA_TTABLE, LUA_TUSERDATA};
use crate::value::{RAW_TVALUE_NIL, RawTValue, TValue};

use super::{LuaInlineDestructor, USERDATA_TAG_IDTOR, Userdata};

const USERDATA_REGISTRY_MEMCAT: u8 = 0;

#[derive(Clone, Copy)]
pub struct UserdataTypeRegistration {
    type_id: TypeId,
    metatable: Table,
}

#[repr(C)]
struct UserdataTypeEntry {
    type_id: TypeId,
    metatable: *mut RawLuaTable,
    destroy: unsafe fn(&Thread, Userdata, Table) -> Result<(), TypedUserdataError>,
    // Only the current registration is a registry root. Older metatables are
    // recognized while existing userdata keep them alive.
    current: bool,
}

#[repr(C)]
pub(crate) struct UserdataTypeRegistry {
    entries: *mut UserdataTypeEntry,
    len: usize,
    cap: usize,
    destructed_metatable: *mut RawLuaTable,
}

impl UserdataTypeEntry {
    #[inline]
    unsafe fn registration(&self) -> UserdataTypeRegistration {
        unsafe {
            UserdataTypeRegistration {
                type_id: self.type_id,
                metatable: Table::from_raw(NonNull::new_unchecked(self.metatable)),
            }
        }
    }
}

impl UserdataTypeRegistry {
    #[inline]
    pub(crate) const fn new() -> Self {
        Self {
            entries: ptr::null_mut(),
            len: 0,
            cap: 0,
            destructed_metatable: ptr::null_mut(),
        }
    }

    #[inline]
    pub(crate) fn get<T: 'static>(&self) -> Option<UserdataTypeRegistration> {
        self.entries()
            .iter()
            .find(|entry| entry.current && entry.type_id == TypeId::of::<T>())
            .map(|entry| unsafe { entry.registration() })
    }

    fn typed_userdata(&self, userdata: Userdata) -> Option<TypedUserdata> {
        unsafe {
            if userdata.as_ptr().as_ref().unwrap_unchecked().tag as usize != USERDATA_TAG_IDTOR {
                return None;
            }

            let metatable = userdata.metatable()?;
            if self.destructed_metatable == metatable.as_ptr() {
                return Some(TypedUserdata {
                    userdata,
                    type_id: None,
                });
            }
            let type_id = self
                .entries()
                .iter()
                .find(|entry| entry.metatable == metatable.as_ptr())?
                .type_id;
            Some(TypedUserdata {
                userdata,
                type_id: Some(type_id),
            })
        }
    }

    pub(crate) unsafe fn register<T: 'static>(
        registry: *mut Self,
        thread: &Thread,
        metatable: Table,
    ) -> VmErrorResult<UserdataTypeRegistration> {
        unsafe {
            if (*registry).len == (*registry).cap {
                Self::grow(registry, thread)?;
            }

            for entry in (*registry).entries_mut() {
                if entry.type_id == TypeId::of::<T>() {
                    entry.current = false;
                }
            }

            let entry = UserdataTypeEntry {
                type_id: TypeId::of::<T>(),
                metatable: metatable.as_ptr(),
                destroy: destroy_typed_userdata_payload::<T>,
                current: true,
            };
            let entry_slot = (*registry).entries.add((*registry).len);
            entry_slot.write(entry);
            (*registry).len += 1;
            Ok((*entry_slot).registration())
        }
    }

    pub(crate) fn retire<T: 'static>(&mut self) {
        let type_id = TypeId::of::<T>();
        for entry in self.entries_mut() {
            if entry.type_id == type_id {
                entry.current = false;
            }
        }
    }

    #[inline]
    pub(crate) fn rooted_metatables(&self) -> impl Iterator<Item = Table> + '_ {
        self.entries()
            .iter()
            .filter(|entry| entry.current)
            .filter_map(|entry| unsafe {
                NonNull::new(entry.metatable).map(|raw| Table::from_raw(raw))
            })
            .chain(
                NonNull::new(self.destructed_metatable).map(|raw| unsafe { Table::from_raw(raw) }),
            )
    }

    #[inline]
    pub(crate) fn recognized_metatables(&self) -> impl Iterator<Item = Table> + '_ {
        self.entries().iter().filter_map(|entry| unsafe {
            NonNull::new(entry.metatable).map(|raw| Table::from_raw(raw))
        })
    }

    pub(crate) unsafe fn remove_unreachable_metatables(&mut self) {
        let mut destination = 0;

        for source in 0..self.len {
            unsafe {
                let entry = self.entries.add(source);
                let metatable = Table::from_raw(NonNull::new_unchecked((*entry).metatable));
                // Atomic marking has reached every live userdata and therefore
                // every historical metatable that still identifies one.
                if !(*entry).current && GcObject::from(metatable).is_white() {
                    ptr::drop_in_place(entry);
                    continue;
                }

                if destination != source {
                    self.entries.add(destination).write(entry.read());
                }
                destination += 1;
            }
        }

        self.len = destination;
    }

    unsafe fn ensure_destructed_metatable(registry: *mut Self, thread: &Thread) -> VmErrorResult {
        if unsafe { !(*registry).destructed_metatable.is_null() } {
            return Ok(());
        }

        unsafe {
            let _stack = StackGuard::new(thread);
            thread.create_table(0, 24)?;
            let metatable = thread.to_object(-1).unwrap_unchecked().table_value();

            thread.push_native_function(destructed_userdata_error, Some("destructed userdata"))?;
            for name in [
                "__add",
                "__sub",
                "__mul",
                "__div",
                "__idiv",
                "__mod",
                "__pow",
                "__unm",
                "__eq",
                "__lt",
                "__le",
                "__len",
                "__concat",
                "__index",
                "__newindex",
                "__call",
                "__tostring",
                "__todebugstring",
                "__namecall",
                "__iter",
            ] {
                thread.push_value(-1)?;
                thread.raw_set_field(-3, name)?;
            }
            thread.pop(1);
            thread.push_boolean(0)?;
            thread.raw_set_field(-2, "__metatable")?;
            thread.set_readonly(-1, 1);
            (*registry).destructed_metatable = metatable.as_ptr();
            Ok(())
        }
    }

    fn destroy(&self, thread: &Thread, userdata: TypedUserdata) -> Result<(), TypedUserdataError> {
        let type_id = userdata.type_id.ok_or(TypedUserdataError::Destructed)?;
        let entry = self
            .entries()
            .iter()
            .find(|entry| entry.type_id == type_id)
            .ok_or(TypedUserdataError::TypeMismatch)?;
        let metatable =
            unsafe { Table::from_raw(NonNull::new_unchecked(self.destructed_metatable)) };
        unsafe { (entry.destroy)(thread, userdata.userdata, metatable) }
    }

    pub(crate) unsafe fn free(registry: *mut Self, thread: &Thread) {
        let entries = unsafe { (*registry).entries };
        let len = unsafe { (*registry).len };
        let cap = unsafe { (*registry).cap };
        unsafe { *registry = Self::new() };
        if entries.is_null() {
            return;
        }

        unsafe {
            for entry in slice::from_raw_parts_mut(entries, len) {
                ptr::drop_in_place(entry);
            }
            thread.free_array(entries, cap, USERDATA_REGISTRY_MEMCAT);
        }
    }

    unsafe fn grow(registry: *mut Self, thread: &Thread) -> VmErrorResult {
        let cap = unsafe { (*registry).cap };
        let new_cap = if cap == 0 {
            4
        } else {
            let Some(new_cap) = cap.checked_mul(2) else {
                return unsafe { thread.too_big() };
            };
            new_cap
        };

        unsafe {
            let entries = thread.realloc_array(
                (*registry).entries,
                cap,
                new_cap,
                USERDATA_REGISTRY_MEMCAT,
            )?;
            (*registry).entries = entries;
            (*registry).cap = new_cap;
        }

        Ok(())
    }

    #[inline]
    fn entries(&self) -> &[UserdataTypeEntry] {
        if self.len == 0 {
            &[]
        } else {
            unsafe { slice::from_raw_parts(self.entries, self.len) }
        }
    }

    #[inline]
    fn entries_mut(&mut self) -> &mut [UserdataTypeEntry] {
        if self.len == 0 {
            &mut []
        } else {
            unsafe { slice::from_raw_parts_mut(self.entries, self.len) }
        }
    }
}

impl GlobalState {
    pub(crate) fn userdata_type_registry_ptr(&self) -> *mut UserdataTypeRegistry {
        unsafe {
            self.as_ptr()
                .as_ref()
                .unwrap_unchecked()
                .userdata_type_registry
                .get()
        }
    }
}

#[derive(Clone, Copy, PartialEq, Eq)]
pub struct TypedUserdata {
    userdata: Userdata,
    type_id: Option<TypeId>,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TypedUserdataError {
    TypeMismatch,
    Destructed,
    Borrowed,
}

/// Unstable registration capability for Rust values stored in VM userdata.
///
/// Registered types are available to VM-native callbacks and to higher-level
/// embedding layers. Raw and tagged userdata remain independent of this
/// registry.
///
/// # Safety
///
/// The thread must be live, stack indices must be valid, registrations must
/// belong to the same VM, and callers must serialize registry mutation and
/// keep referenced metatables rooted. Each registered metatable must identify
/// exactly one Rust payload type and must not be reassigned to userdata with a
/// different payload layout.
#[allow(
    clippy::missing_safety_doc,
    reason = "all methods share the capability-level safety contract"
)]
pub trait UserdataTypeRegistryAccess: Sealed {
    unsafe fn userdata_type<T: 'static>(&self) -> Option<UserdataTypeRegistration>;

    unsafe fn retire_userdata_type<T: 'static>(&self);

    unsafe fn register_userdata_type<T: 'static>(
        &self,
        metatable_index: i32,
    ) -> VmErrorResult<UserdataTypeRegistration>;
}

/// Unstable access capability for registered Rust userdata.
///
/// This provides the storage and validation mechanics used by VM-native
/// callbacks and higher-level embedding layers; it does not root returned
/// userdata descriptors.
///
/// # Safety
///
/// The thread and registration must be live and belong to the same VM. Stack
/// indices must be valid. Returned descriptors do not root their userdata;
/// callers must keep the object rooted and must not replace its registered
/// metatable or replace its payload while a descriptor or borrow remains live.
/// The lifecycle operations require exclusive payload access. They remove the
/// Rust payload but leave the associated value slot intact. Values returned
/// from that slot are non-owning VM views and remain valid only while their
/// userdata and referenced object remain live.
#[allow(
    clippy::missing_safety_doc,
    reason = "all methods share the capability-level safety contract"
)]
pub trait TypedUserdataAccess: Sealed {
    unsafe fn typed_userdata(&self, userdata: Userdata) -> Option<TypedUserdata>;

    unsafe fn typed_userdata_at(&self, index: i32) -> Option<TypedUserdata>;

    unsafe fn push_typed_userdata<T: 'static>(
        &self,
        value: T,
        registration: &UserdataTypeRegistration,
    ) -> VmErrorResult;

    unsafe fn typed_userdata_value(&self, userdata: TypedUserdata) -> TValue;

    unsafe fn set_typed_userdata_value(&self, userdata: TypedUserdata, value: TValue);

    unsafe fn take_typed_userdata<T: 'static>(
        &self,
        userdata: TypedUserdata,
    ) -> Result<T, TypedUserdataError>;

    unsafe fn destroy_typed_userdata(
        &self,
        userdata: TypedUserdata,
    ) -> Result<(), TypedUserdataError>;
}

impl UserdataTypeRegistryAccess for Thread {
    unsafe fn userdata_type<T: 'static>(&self) -> Option<UserdataTypeRegistration> {
        unsafe { (&*self.global().userdata_type_registry_ptr()).get::<T>() }
    }

    unsafe fn retire_userdata_type<T: 'static>(&self) {
        unsafe { (&mut *self.global().userdata_type_registry_ptr()).retire::<T>() }
    }

    unsafe fn register_userdata_type<T: 'static>(
        &self,
        metatable_index: i32,
    ) -> VmErrorResult<UserdataTypeRegistration> {
        unsafe {
            if self.type_of(metatable_index) != LUA_TTABLE {
                return crate::error!(self, "userdata metatable must be a table");
            }

            let metatable = self
                .to_object(metatable_index)
                .unwrap_unchecked()
                .table_value();
            UserdataTypeRegistry::ensure_destructed_metatable(
                self.global().userdata_type_registry_ptr(),
                self,
            )?;
            UserdataTypeRegistry::register::<T>(
                self.global().userdata_type_registry_ptr(),
                self,
                metatable,
            )
        }
    }
}

impl TypedUserdataAccess for Thread {
    unsafe fn typed_userdata(&self, userdata: Userdata) -> Option<TypedUserdata> {
        unsafe { (&*self.global().userdata_type_registry_ptr()).typed_userdata(userdata) }
    }

    unsafe fn typed_userdata_at(&self, index: i32) -> Option<TypedUserdata> {
        unsafe {
            let object = self.to_object(index)?;
            (object.tt() == LUA_TUSERDATA)
                .then(|| object.userdata_value())
                .and_then(|userdata| self.typed_userdata(userdata))
        }
    }

    unsafe fn push_typed_userdata<T: 'static>(
        &self,
        value: T,
        registration: &UserdataTypeRegistration,
    ) -> VmErrorResult {
        unsafe {
            let mut stack = StackGuard::new(self);
            self.lua_check_stack(2, None)?;

            if registration.type_id != TypeId::of::<T>() {
                return crate::error!(self, "userdata registration does not match payload type");
            }

            let Some(payload_len) = typed_userdata_payload_len::<T>() else {
                return self.too_big();
            };
            let data = self.new_userdata_dtor(payload_len, drop_typed_userdata::<T>)?;
            typed_userdata_value_ptr(data.cast()).write(RAW_TVALUE_NIL);
            typed_userdata_cell_ptr::<T>(data.cast()).write(RefCell::new(Some(value)));

            self.thread_barrier();
            self.push_table(registration.metatable)?;
            self.set_metatable(-2)?;
            stack.keep(1);
        }
        Ok(())
    }

    unsafe fn typed_userdata_value(&self, userdata: TypedUserdata) -> TValue {
        unsafe { userdata.associated_value() }
    }

    unsafe fn set_typed_userdata_value(&self, userdata: TypedUserdata, value: TValue) {
        unsafe {
            userdata.associated_value().set_obj(value);
            self.barrier_value(userdata.userdata.into(), value);
        }
    }

    unsafe fn take_typed_userdata<T: 'static>(
        &self,
        userdata: TypedUserdata,
    ) -> Result<T, TypedUserdataError> {
        unsafe {
            if userdata.type_id.is_none() {
                return Err(TypedUserdataError::Destructed);
            }
            let Some(cell) = userdata.cell_ptr::<T>() else {
                return Err(TypedUserdataError::TypeMismatch);
            };
            let mut value = (&*cell)
                .try_borrow_mut()
                .map_err(|_| TypedUserdataError::Borrowed)?;
            let value = value.take().ok_or(TypedUserdataError::Destructed)?;
            let registry = &*self.global().userdata_type_registry_ptr();
            let metatable = Table::from_raw(NonNull::new_unchecked(registry.destructed_metatable));
            mark_userdata_destructed(self, userdata.userdata, metatable);
            Ok(value)
        }
    }

    unsafe fn destroy_typed_userdata(
        &self,
        userdata: TypedUserdata,
    ) -> Result<(), TypedUserdataError> {
        unsafe { (&*self.global().userdata_type_registry_ptr()).destroy(self, userdata) }
    }
}

impl TypedUserdata {
    pub fn is<T: 'static>(&self) -> bool {
        self.type_id == Some(TypeId::of::<T>())
    }

    pub const fn is_destructed(&self) -> bool {
        self.type_id.is_none()
    }

    unsafe fn associated_value(&self) -> TValue {
        unsafe {
            TValue::from_raw(NonNull::new_unchecked(typed_userdata_value_ptr(
                self.userdata.data_mut_ptr(),
            )))
        }
    }

    /// Returns the typed userdata cell address for `T`.
    ///
    /// # Safety
    ///
    /// The userdata object that produced this descriptor must remain alive for
    /// any access through the returned pointer. `TypedUserdata` is not a root.
    pub unsafe fn cell_ptr<T: 'static>(&self) -> Option<*mut RefCell<Option<T>>> {
        if !self.is::<T>() {
            return None;
        }

        let payload_len = self
            .userdata
            .total_payload_len()
            .checked_sub(core::mem::size_of::<LuaInlineDestructor>())?;
        if payload_len < typed_userdata_payload_len::<T>()? {
            return None;
        }

        Some(unsafe { typed_userdata_cell_ptr::<T>(self.userdata.data_mut_ptr()) })
    }
}

fn typed_userdata_payload_len<T>() -> Option<usize> {
    (core::mem::align_of::<RawTValue>() - 1)
        .checked_add(core::mem::size_of::<RawTValue>())?
        .checked_add(core::mem::align_of::<RefCell<Option<T>>>() - 1)?
        .checked_add(core::mem::size_of::<RefCell<Option<T>>>())
}

unsafe fn typed_userdata_value_ptr(data: *mut u8) -> *mut RawTValue {
    let offset = data.align_offset(core::mem::align_of::<RawTValue>());
    debug_assert_ne!(offset, usize::MAX);
    unsafe { data.add(offset).cast() }
}

unsafe fn typed_userdata_cell_ptr<T>(data: *mut u8) -> *mut RefCell<Option<T>> {
    let value_end = unsafe {
        typed_userdata_value_ptr(data)
            .cast::<u8>()
            .add(core::mem::size_of::<RawTValue>())
    };
    let offset = value_end.align_offset(core::mem::align_of::<RefCell<Option<T>>>());
    debug_assert_ne!(offset, usize::MAX);
    unsafe { value_end.add(offset).cast() }
}

fn drop_typed_userdata<T>(data: *mut ()) {
    unsafe {
        ptr::drop_in_place(typed_userdata_cell_ptr::<T>(data.cast()));
    }
}

unsafe fn destroy_typed_userdata_payload<T>(
    thread: &Thread,
    userdata: Userdata,
    metatable: Table,
) -> Result<(), TypedUserdataError> {
    unsafe {
        let cell = &*typed_userdata_cell_ptr::<T>(userdata.data_mut_ptr());
        let mut value = cell
            .try_borrow_mut()
            .map_err(|_| TypedUserdataError::Borrowed)?;
        let value = value.take().ok_or(TypedUserdataError::Destructed)?;
        mark_userdata_destructed(thread, userdata, metatable);
        drop(value);
        Ok(())
    }
}

unsafe fn mark_userdata_destructed(thread: &Thread, userdata: Userdata, metatable: Table) {
    unsafe {
        userdata.set_metatable(Some(metatable));
        thread.object_barrier(userdata.into(), metatable.into());
    }
}

fn destructed_userdata_error(context: NativeCallContext<'_>) -> NativeCallResult {
    unsafe {
        crate::error!(context.raw_thread(), "userdata has been destructed").map_err(Into::into)
    }
}