luau-vm 0.732.0

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

use luau_common::BStr;
use luau_printf::Arg;

use crate::thread::Thread;
use crate::userdata::{TypedUserdataAccess, UserdataTypeRegistryAccess};
use crate::{VmErrorResult, VmResult};

pub type NativeCallResult = VmResult<usize>;

pub type RawNativeFunction = for<'call> fn(NativeCallContext<'call>) -> NativeCallResult;

pub type RawNativeContinuation = for<'call> fn(NativeCallContext<'call>, i32) -> NativeCallResult;

// Native function registration
#[derive(Clone, Copy)]
pub struct NativeFunction {
    pub name: &'static str,
    pub function: RawNativeFunction,
}

#[derive(Clone, Copy)]
pub struct NativeModule {
    pub name: Option<&'static str>,
    pub functions: &'static [NativeFunction],
}

impl NativeModule {
    #[inline]
    pub const fn new(name: Option<&'static str>, functions: &'static [NativeFunction]) -> Self {
        Self { name, functions }
    }
}

// Native call context
#[derive(PartialEq, Eq)]
#[repr(transparent)]
pub struct NativeCallContext<'call> {
    thread: &'call Thread,
}

impl<'call> NativeCallContext<'call> {
    /// The thread must be the active thread for a VM-native call or continuation.
    #[inline]
    pub(crate) const fn new(thread: &'call Thread) -> Self {
        Self { thread }
    }

    #[inline]
    pub const fn raw_thread(&self) -> &'call Thread {
        self.thread
    }

    #[inline]
    pub fn arg_count(&self) -> i32 {
        unsafe { self.thread.get_top() }
    }

    #[inline]
    pub fn top(&self) -> i32 {
        self.arg_count()
    }

    #[inline]
    pub fn pop(&self, count: i32) {
        unsafe { self.thread.pop(count) }
    }

    #[inline]
    pub fn arg(&self, index: i32) -> NativeArgument<'_, 'call> {
        NativeArgument {
            context: self,
            index,
        }
    }

    #[inline]
    pub fn argument(&self, index: i32) -> NativeArgument<'_, 'call> {
        self.arg(index)
    }

    #[inline]
    pub fn upvalue(&self, index: i32) -> NativeArgument<'_, 'call> {
        NativeArgument {
            context: self,
            index: crate::thread::upvalue_index(index),
        }
    }

    #[inline]
    pub fn args(&self) -> NativeArguments<'_, 'call> {
        NativeArguments {
            context: self,
            next: 1,
            end: self.arg_count(),
        }
    }

    #[inline]
    pub fn push_number(&self, value: f64) -> VmErrorResult {
        unsafe { self.thread.push_number(value) }
    }

    #[inline]
    pub fn push_nil(&self) -> VmErrorResult {
        unsafe { self.thread.push_nil() }
    }

    #[inline]
    pub fn push_integer(&self, value: i32) -> VmErrorResult {
        unsafe { self.thread.push_integer(value) }
    }

    #[inline]
    pub fn push_integer64(&self, value: i64) -> VmErrorResult {
        unsafe { self.thread.push_integer64(value) }
    }

    #[inline]
    pub fn push_unsigned(&self, value: u32) -> VmErrorResult {
        unsafe { self.thread.push_unsigned(value) }
    }

    #[inline]
    pub fn push_vector(&self, components: [f32; crate::types::LUA_VECTOR_SIZE]) -> VmErrorResult {
        unsafe { self.thread.push_vector(components) }
    }

    #[inline]
    pub fn push_boolean(&self, value: bool) -> VmErrorResult {
        unsafe { self.thread.push_boolean(i32::from(value)) }
    }

    #[inline]
    pub fn push_bool(&self, value: bool) -> VmErrorResult {
        self.push_boolean(value)
    }

    #[inline]
    pub fn push_string(&self, bytes: impl AsRef<[u8]>) -> VmErrorResult {
        unsafe { self.thread.push_string(bytes) }
    }

    #[inline]
    pub fn push_userdata<T: 'static>(&self, value: T) -> VmErrorResult {
        let Some(registration) = (unsafe { self.thread.userdata_type::<T>() }) else {
            return self.error("userdata type is not registered", []);
        };
        unsafe { self.thread.push_typed_userdata(value, &registration) }
    }

    #[inline]
    pub fn lua_error<'a, T>(
        &self,
        format: impl AsRef<[u8]>,
        args: impl AsMut<[Arg<'a>]>,
    ) -> VmErrorResult<T> {
        unsafe { self.thread.lua_error(format, args) }
    }

    #[inline]
    pub fn error<'a, T>(
        &self,
        format: impl AsRef<[u8]>,
        args: impl AsMut<[Arg<'a>]>,
    ) -> VmErrorResult<T> {
        self.lua_error(format, args)
    }
}

// Native arguments
#[derive(Clone, Copy)]
pub struct NativeArgument<'ctx, 'call> {
    pub(super) context: &'ctx NativeCallContext<'call>,
    pub(super) index: i32,
}

impl<'ctx, 'call> NativeArgument<'ctx, 'call> {
    #[inline]
    pub const fn index(&self) -> i32 {
        self.index
    }

    #[inline]
    pub fn integer(&self) -> VmErrorResult<i32> {
        unsafe { self.context.raw_thread().check_integer(self.index) }
    }

    #[inline]
    pub fn integer_or(&self, default: i32) -> VmErrorResult<i32> {
        unsafe { self.context.raw_thread().opt_integer(self.index, default) }
    }

    #[inline]
    pub fn number(&self) -> VmErrorResult<f64> {
        unsafe { self.context.raw_thread().check_number(self.index) }
    }

    #[inline]
    pub fn number_or(&self, default: f64) -> VmErrorResult<f64> {
        unsafe { self.context.raw_thread().opt_number(self.index, default) }
    }

    #[inline]
    pub fn integer64(&self) -> VmErrorResult<i64> {
        unsafe { self.context.raw_thread().check_integer64(self.index) }
    }

    #[inline]
    pub fn integer64_or(&self, default: i64) -> VmErrorResult<i64> {
        unsafe { self.context.raw_thread().opt_integer64(self.index, default) }
    }

    #[inline]
    pub fn unsigned(&self) -> VmErrorResult<u32> {
        unsafe { self.context.raw_thread().check_unsigned(self.index) }
    }

    #[inline]
    pub fn vector(&self) -> VmErrorResult<[f32; crate::types::LUA_VECTOR_SIZE]> {
        unsafe { self.context.raw_thread().check_vector(self.index) }
    }

    #[inline]
    /// Returns a borrowed string argument.
    ///
    /// # Safety
    ///
    /// The argument must remain rooted and no VM transition may pop, unroot,
    /// or collect it for the returned borrow's lifetime.
    pub unsafe fn string(&self) -> VmErrorResult<&'ctx BStr> {
        unsafe { self.context.raw_thread().check_string(self.index) }
    }

    #[inline]
    pub fn light_userdata(&self) -> VmErrorResult<*mut ()> {
        let pointer = unsafe { self.context.raw_thread().to_light_userdata(self.index) };
        if pointer.is_null() {
            return self.type_error("light userdata");
        }
        Ok(pointer)
    }

    /// Borrows this argument as userdata of the registered type.
    #[inline]
    /// # Safety
    ///
    /// The argument must remain rooted and no VM transition may pop, unroot,
    /// or collect it while the returned guard is alive.
    pub unsafe fn userdata<T: 'static>(&self) -> VmErrorResult<LuaUserdataRef<'ctx, T>> {
        let cell = unsafe { &*self.userdata_cell_ptr::<T>()? };
        let Ok(value) = cell.try_borrow() else {
            return self.error("userdata is already mutably borrowed");
        };
        let Ok(value) = Ref::filter_map(value, Option::as_ref) else {
            return self.error("userdata has been destructed");
        };
        Ok(LuaUserdataRef::new(value))
    }

    /// Mutably borrows this argument as userdata of the registered type.
    #[inline]
    /// # Safety
    ///
    /// The argument must remain rooted and no VM transition may pop, unroot,
    /// or collect it while the returned guard is alive. No other access to the
    /// same typed value may overlap the mutable guard.
    pub unsafe fn userdata_mut<T: 'static>(&self) -> VmErrorResult<LuaUserdataRefMut<'ctx, T>> {
        let cell = unsafe { &*self.userdata_cell_ptr::<T>()? };
        let Ok(value) = cell.try_borrow_mut() else {
            return self.error("userdata is already borrowed");
        };
        let Ok(value) = RefMut::filter_map(value, Option::as_mut) else {
            return self.error("userdata has been destructed");
        };
        Ok(LuaUserdataRefMut::new(value))
    }

    #[inline]
    unsafe fn userdata_cell_ptr<T: 'static>(&self) -> VmErrorResult<*mut RefCell<Option<T>>> {
        let expected_type = core::any::type_name::<T>();
        let Some(userdata) = (unsafe { self.context.raw_thread().typed_userdata_at(self.index) })
        else {
            return self.type_error(expected_type);
        };
        if userdata.is_destructed() {
            return self.error("userdata has been destructed");
        }
        let Some(cell) = (unsafe { userdata.cell_ptr::<T>() }) else {
            return self.type_error(expected_type);
        };

        Ok(cell)
    }

    #[inline]
    pub fn error<T>(&self, message: impl AsRef<[u8]>) -> VmErrorResult<T> {
        unsafe { self.context.raw_thread().lua_arg_error(self.index, message) }
    }

    #[inline]
    pub fn expected(&self, condition: bool, expected_type: &str) -> VmErrorResult {
        unsafe {
            self.context
                .raw_thread()
                .lua_arg_expected(condition, self.index, expected_type)
        }
    }

    #[inline]
    pub fn type_error<T>(&self, expected_type: &str) -> VmErrorResult<T> {
        unsafe {
            self.context
                .raw_thread()
                .lua_type_error(self.index, expected_type)
        }
    }
}

#[derive(Clone, Copy)]
pub struct NativeArguments<'ctx, 'call> {
    pub(super) context: &'ctx NativeCallContext<'call>,
    pub(super) next: i32,
    pub(super) end: i32,
}

impl NativeArguments<'_, '_> {
    #[inline]
    pub fn remaining(&self) -> i32 {
        (self.end - self.next + 1).max(0)
    }
}

impl<'ctx, 'call> Iterator for NativeArguments<'ctx, 'call> {
    type Item = NativeArgument<'ctx, 'call>;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if self.next > self.end {
            None
        } else {
            let index = self.next;
            self.next += 1;
            Some(NativeArgument {
                context: self.context,
                index,
            })
        }
    }

    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        let remaining = self.remaining() as usize;
        (remaining, Some(remaining))
    }
}

impl ExactSizeIterator for NativeArguments<'_, '_> {}

/// Shared borrow of a Lua-owned userdata value.
///
/// The borrow is released when this guard is dropped.
pub struct LuaUserdataRef<'lua, T> {
    value: Ref<'lua, T>,
}

impl<'lua, T> LuaUserdataRef<'lua, T> {
    #[inline]
    pub(super) fn new(value: Ref<'lua, T>) -> Self {
        Self { value }
    }
}

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

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

/// Mutable borrow of a Lua-owned userdata value.
///
/// The borrow is released when this guard is dropped.
pub struct LuaUserdataRefMut<'lua, T> {
    value: RefMut<'lua, T>,
}

impl<'lua, T> LuaUserdataRefMut<'lua, T> {
    #[inline]
    pub(super) fn new(value: RefMut<'lua, T>) -> Self {
        Self { value }
    }
}

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

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

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