luau-vm 0.732.0

use core::mem::{offset_of, size_of};
use core::ptr::{self, NonNull};

use crate::Table;
use crate::gc::GcObject;
use crate::gc::RawGcObject;
use crate::handle::RawHandle;
use crate::native::{RawNativeContinuation, RawNativeFunction};
use crate::string::{LuaString, RawTString, TString};
use crate::table::RawLuaTable;
use crate::value::{RawTValue, TValue, TValueCursor};

use super::{Proto, RawProto};

#[repr(C)]
pub struct RawClosure {
    pub tt: u8,
    pub marked: u8,
    pub memcat: u8,
    pub is_native: u8,
    pub n_upvalues: u8,
    pub stack_size: u8,
    pub preload: u8,
    pub gc_list: *mut RawGcObject,
    pub env: *mut RawLuaTable,
    pub data: [u8; 0],
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct RawNativeClosure {
    pub function: Option<RawNativeFunction>,
    pub continuation: Option<RawNativeContinuation>,
    pub debug_name_deprecated: Option<&'static str>,
    pub debug_name: *mut RawTString,
}

#[repr(C)]
#[derive(Clone, Copy)]
pub struct RawLuaClosure {
    pub proto: *mut RawProto,
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
/// Non-owning identity of a VM closure record.
///
/// # Safety model for unsafe methods
///
/// The closure and referenced environment, proto, upvalues, and GC objects
/// must remain live in the same VM. The native/Luau variant and every index
/// must match the allocated closure layout.
pub struct Closure {
    raw: NonNull<RawClosure>,
}

#[allow(
    clippy::missing_safety_doc,
    reason = "Closure's shared raw-handle contract is documented on Closure"
)]
impl Closure {
    pub const unsafe fn from_raw(raw: NonNull<RawClosure>) -> Self {
        Self { raw }
    }

    pub fn size_native_closure(n: usize) -> usize {
        offset_of!(RawClosure, data) + size_of::<RawNativeClosure>() + size_of::<RawTValue>() * n
    }

    pub fn size_lua_closure(n: usize) -> usize {
        offset_of!(RawClosure, data) + size_of::<RawLuaClosure>() + size_of::<RawTValue>() * n
    }

    /// `sizeclosure`
    pub unsafe fn size(&self) -> usize {
        unsafe {
            if self.is_native() {
                Closure::size_native_closure(
                    self.as_ptr().as_ref().unwrap_unchecked().n_upvalues as usize,
                )
            } else {
                Closure::size_lua_closure(
                    self.as_ptr().as_ref().unwrap_unchecked().n_upvalues as usize,
                )
            }
        }
    }

    pub unsafe fn is_native(&self) -> bool {
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().is_native != 0 }
    }

    pub unsafe fn is_lua(&self) -> bool {
        unsafe { !self.is_native() }
    }

    pub unsafe fn proto(&self) -> Option<Proto> {
        unsafe {
            if self.is_lua() {
                NonNull::new(self.lua_data().proto).map(|raw| Proto::from_raw(raw))
            } else {
                None
            }
        }
    }

    pub unsafe fn env(&self) -> Table {
        unsafe {
            Table::from_raw(NonNull::new_unchecked(
                self.as_ptr().as_ref().unwrap_unchecked().env,
            ))
        }
    }

    pub unsafe fn gc_list(&self) -> Option<GcObject> {
        unsafe {
            NonNull::new(self.as_ptr().as_ref().unwrap_unchecked().gc_list)
                .map(|raw| GcObject::from_raw(raw))
        }
    }

    pub unsafe fn set_gc_list(&self, gc_list: Option<GcObject>) {
        unsafe {
            self.as_ptr().as_mut().unwrap_unchecked().gc_list =
                gc_list.map_or(ptr::null_mut(), |object| object.as_ptr());
        }
    }

    pub unsafe fn set_env(&self, env: Table) {
        unsafe {
            self.as_ptr().as_mut().unwrap_unchecked().env = env.as_ptr();
        }
    }

    pub unsafe fn native_data(&self) -> RawNativeClosure {
        debug_assert!(unsafe { self.is_native() });
        unsafe { self.data_ptr().cast::<RawNativeClosure>().read() }
    }

    pub unsafe fn native_data_ptr(&self) -> *mut RawNativeClosure {
        debug_assert!(unsafe { self.is_native() });
        self.data_ptr().cast()
    }

    pub unsafe fn set_native_data(&self, data: RawNativeClosure) {
        unsafe { self.native_data_ptr().write(data) };
    }

    pub unsafe fn native_debug_name(&self) -> Option<LuaString> {
        debug_assert!(unsafe { self.is_native() });
        let data = unsafe { self.native_data() };
        if luau_common::flags::LuauManagedDebugNames.get() {
            NonNull::new(data.debug_name)
                .map(|raw| LuaString::from_interned(unsafe { TString::from_raw(raw) }))
        } else {
            data.debug_name_deprecated
                .map(|name| LuaString::from_static(luau_common::BStr::new(name.as_bytes())))
        }
    }

    pub unsafe fn lua_data(&self) -> RawLuaClosure {
        debug_assert!(unsafe { self.is_lua() });
        unsafe { self.data_ptr().cast::<RawLuaClosure>().read() }
    }

    pub unsafe fn lua_data_ptr(&self) -> *mut RawLuaClosure {
        debug_assert!(unsafe { self.is_lua() });
        self.data_ptr().cast()
    }

    pub unsafe fn set_lua_proto(&self, proto: Proto) {
        unsafe {
            self.lua_data_ptr().write(RawLuaClosure {
                proto: proto.as_ptr(),
            })
        };
    }

    pub unsafe fn native_upvalue(&self, index: usize) -> TValue {
        unsafe {
            debug_assert!(index < self.as_ptr().as_ref().unwrap_unchecked().n_upvalues as usize);
            TValue::from_raw(NonNull::new_unchecked(
                self.native_upvalues_ptr().add(index),
            ))
        }
    }

    pub unsafe fn native_upvalues(&self) -> TValueCursor {
        unsafe { TValueCursor::from_ptr(self.native_upvalues_ptr()) }
    }

    pub unsafe fn lua_upvalue_ref(&self, index: usize) -> TValue {
        unsafe {
            debug_assert!(index < self.as_ptr().as_ref().unwrap_unchecked().n_upvalues as usize);
            TValue::from_raw(NonNull::new_unchecked(
                self.lua_upvalue_refs_ptr().add(index),
            ))
        }
    }

    pub unsafe fn lua_upvalue_refs(&self) -> TValueCursor {
        unsafe { TValueCursor::from_ptr(self.lua_upvalue_refs_ptr()) }
    }

    fn data_ptr(&self) -> *mut u8 {
        unsafe { self.as_ptr().cast::<u8>().add(offset_of!(RawClosure, data)) }
    }

    unsafe fn native_upvalues_ptr(&self) -> *mut RawTValue {
        debug_assert!(unsafe { self.is_native() });
        unsafe {
            self.data_ptr()
                .add(size_of::<RawNativeClosure>())
                .cast::<RawTValue>()
        }
    }

    unsafe fn lua_upvalue_refs_ptr(&self) -> *mut RawTValue {
        debug_assert!(unsafe { self.is_lua() });
        unsafe {
            self.data_ptr()
                .add(size_of::<RawLuaClosure>())
                .cast::<RawTValue>()
        }
    }
}
impl crate::handle::sealed::Sealed for Closure {}
impl RawHandle for Closure {
    type Raw = RawClosure;

    fn as_ptr(&self) -> *mut Self::Raw {
        self.raw.as_ptr()
    }
}

impl AsRef<Closure> for Closure {
    fn as_ref(&self) -> &Closure {
        self
    }
}