luau-vm 0.732.0

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

use crate::Table;
use crate::gc::GcObject;
use crate::handle::RawHandle;
use crate::memory::MemoryRuntime;
use crate::string::TString;
use crate::thread::Thread;
use crate::types::{
    LUA_EXTRA_SIZE, LUA_TBOOLEAN, LUA_TDEADKEY, LUA_TINTEGER, LUA_TLIGHTUSERDATA, LUA_TNIL,
    LUA_TNUMBER, LUA_TSTRING, LUA_TVECTOR, LUA_VECTOR_SIZE,
};
use crate::value::{RAW_TVALUE_NIL, RawTValue, RawValue, TValue, TValueCursor};

use super::{RawLuaTable, RawLuaTableFree};

#[derive(Clone, Copy)]
#[repr(C)]
pub struct RawTKey {
    pub value: RawValue,
    pub extra: [i32; LUA_EXTRA_SIZE],
    pub tt_next: u32,
}

pub const RAW_TKEY_NIL: RawTKey = RawTKey {
    value: RawValue {
        pointer: core::ptr::null_mut(),
    },
    extra: [0; LUA_EXTRA_SIZE],
    tt_next: LUA_TNIL as u32,
};

pub const RAW_TKEY_DEAD_KEY: RawTKey = RawTKey {
    value: RawValue {
        pointer: core::ptr::null_mut(),
    },
    extra: [0; LUA_EXTRA_SIZE],
    tt_next: LUA_TDEADKEY as u32,
};

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
/// Non-owning view of a live table-key record.
///
/// Unsafe operations require the owning table storage to remain live and the
/// key's tag/payload and collision-chain encoding to remain valid.
pub struct TKey {
    raw: NonNull<RawTKey>,
}

#[allow(
    clippy::missing_safety_doc,
    reason = "TKey is a non-owning table-key view governed by internal's raw-view contract"
)]
impl TKey {
    const TT_BITS: u32 = 4;
    const TT_MASK: u32 = (1 << Self::TT_BITS) - 1;
    const NEXT_SHIFT: u32 = Self::TT_BITS;
    const NEXT_BITS: u32 = 32 - Self::NEXT_SHIFT;
    const NEXT_MASK: u32 = !Self::TT_MASK;

    pub const unsafe fn from_raw(raw: NonNull<RawTKey>) -> Self {
        Self { raw }
    }

    pub fn set_nil(&self) {
        unsafe {
            *self.as_ptr().as_mut().unwrap_unchecked() = RAW_TKEY_NIL;
        }
    }

    pub fn tt(&self) -> i32 {
        (unsafe { self.as_ptr().as_ref().unwrap_unchecked().tt_next } & Self::TT_MASK) as i32
    }

    /// `setttype`
    pub fn set_tt(&self, tt: i32) {
        unsafe {
            let raw = self.as_ptr().as_mut().unwrap_unchecked();
            raw.tt_next = (raw.tt_next & Self::NEXT_MASK) | (tt as u32 & Self::TT_MASK);
        }
    }

    pub fn next(&self) -> i32 {
        let raw = unsafe { self.as_ptr().as_ref().unwrap_unchecked().tt_next } >> Self::NEXT_SHIFT;
        let shift = 32 - Self::NEXT_BITS;
        ((raw << shift) as i32) >> shift
    }

    /// `gnext`
    pub fn set_next(&self, next: i32) {
        let next_bits = ((next as u32) << Self::NEXT_SHIFT) & Self::NEXT_MASK;
        unsafe {
            let raw = self.as_ptr().as_mut().unwrap_unchecked();
            raw.tt_next = (raw.tt_next & Self::TT_MASK) | next_bits;
        }
    }

    pub fn is_nil(&self) -> bool {
        self.tt() == LUA_TNIL
    }

    pub fn is_dead_key(&self) -> bool {
        self.tt() == LUA_TDEADKEY
    }

    /// `iscollectable`
    pub fn is_collectable(&self) -> bool {
        self.tt() >= LUA_TSTRING
    }

    /// `gcvalue`
    pub fn gc_value(&self) -> GcObject {
        debug_assert!(self.is_collectable());
        unsafe {
            GcObject::from_raw(NonNull::new_unchecked(
                self.as_ptr().as_ref().unwrap_unchecked().value.gc,
            ))
        }
    }

    /// `pvalue`
    pub fn pointer_value(&self) -> *mut () {
        debug_assert!(self.tt() == LUA_TLIGHTUSERDATA);
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().value.pointer }
    }

    /// `nvalue`
    pub fn number_value(&self) -> f64 {
        debug_assert!(self.tt() == LUA_TNUMBER);
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().value.number }
    }

    /// `lvalue`
    pub fn integer_value(&self) -> i64 {
        debug_assert!(self.tt() == LUA_TINTEGER);
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().value.integer }
    }

    /// `bvalue`
    pub fn boolean_value(&self) -> i32 {
        debug_assert!(self.tt() == LUA_TBOOLEAN);
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().value.boolean }
    }

    /// `lightuserdatatag`
    pub fn light_userdata_tag(&self) -> i32 {
        debug_assert!(self.tt() == LUA_TLIGHTUSERDATA);
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().extra[0] }
    }

    /// `vvalue`
    pub fn vector_value(&self) -> [f32; LUA_VECTOR_SIZE] {
        debug_assert!(self.tt() == LUA_TVECTOR);
        let vector = self.as_ptr().cast::<f32>();
        #[cfg(not(feature = "vector4"))]
        unsafe {
            [*vector, *vector.add(1), *vector.add(2)]
        }
        #[cfg(feature = "vector4")]
        unsafe {
            [*vector, *vector.add(1), *vector.add(2), *vector.add(3)]
        }
    }

    /// `luaO_rawequalKey`
    pub fn raw_equal_value(&self, other: impl Into<TValue>) -> bool {
        let other = other.into();
        if self.tt() != unsafe { other.as_ptr().as_ref().unwrap_unchecked().tt } {
            return false;
        }

        match self.tt() {
            x if x == LUA_TNIL => true,
            x if x == LUA_TNUMBER => self.number_value() == other.number_value(),
            x if x == LUA_TINTEGER => self.integer_value() == other.integer_value(),
            x if x == LUA_TVECTOR => self.vector_value() == other.vector_value(),
            x if x == LUA_TBOOLEAN => self.boolean_value() == other.boolean_value(),
            x if x == LUA_TLIGHTUSERDATA => {
                self.pointer_value() == other.pointer_value()
                    && self.light_userdata_tag() == other.light_userdata_tag()
            }
            _ => {
                debug_assert!(self.is_collectable());
                self.gc_value() == other.gc_value()
            }
        }
    }
}
impl crate::handle::sealed::Sealed for TKey {}
impl RawHandle for TKey {
    type Raw = RawTKey;

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

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

#[repr(C)]
pub struct RawLuaNode {
    pub value: RawTValue,
    pub key: RawTKey,
}

pub const RAW_LUA_NODE_DUMMY: RawLuaNode = RawLuaNode {
    value: RAW_TVALUE_NIL,
    key: RAW_TKEY_NIL,
};

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
/// Non-owning view of a live table hash node.
///
/// Unsafe operations require the owning table's node array to remain live and
/// the node's key/value records to remain valid.
pub struct LuaNode {
    raw: NonNull<RawLuaNode>,
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
/// Nullable traversal position in a table's hash-node array.
///
/// Unsafe navigation requires cursors from the same live node allocation and
/// an in-bounds non-null position where a node is accessed. Resizing the table
/// invalidates every cursor.
pub struct LuaNodeCursor(*mut RawLuaNode);

#[allow(
    clippy::missing_safety_doc,
    reason = "LuaNode is a non-owning table-node view governed by Table's contract"
)]
impl LuaNode {
    pub const unsafe fn from_raw(raw: NonNull<RawLuaNode>) -> Self {
        Self { raw }
    }

    pub fn value(&self) -> TValue {
        unsafe { TValue::from_raw(NonNull::new_unchecked(&raw mut (*self.as_ptr()).value)) }
    }

    pub fn value_unchecked(&self) -> TValue {
        self.value()
    }

    pub fn key(&self) -> TKey {
        unsafe { TKey::from_raw(NonNull::new_unchecked(&raw mut (*self.as_ptr()).key)) }
    }

    pub fn has_string_key(&self, key: TString) -> bool {
        unsafe {
            const TT_MASK: u32 = (1 << 4) - 1;

            let raw = self.as_ptr();
            let key_raw = &raw const (*raw).key;
            ((*key_raw).tt_next & TT_MASK) as i32 == LUA_TSTRING
                && (*key_raw).value.gc.cast() == key.as_ptr()
        }
    }

    pub fn value_is_nil(&self) -> bool {
        unsafe { (*self.as_ptr()).value.tt == LUA_TNIL }
    }

    /// `gnext`
    pub fn next(&self) -> i32 {
        self.key().next()
    }

    /// `setnodekey`
    pub fn set_key_from_value(&self, value: impl Into<TValue>) {
        let value = value.into();
        unsafe {
            ptr::copy_nonoverlapping(
                (&raw const (*value.as_ptr()).value).cast::<u8>(),
                (&raw mut (*self.as_ptr()).key.value).cast::<u8>(),
                core::mem::size_of::<RawValue>(),
            );
            ptr::copy_nonoverlapping(
                (&raw const (*value.as_ptr()).extra).cast::<u8>(),
                (&raw mut (*self.as_ptr()).key.extra).cast::<u8>(),
                core::mem::size_of::<[i32; LUA_EXTRA_SIZE]>(),
            );
            self.key().set_tt((*value.as_ptr()).tt);
        }
    }

    /// `getnodekey`
    pub fn write_key_to_value(&self, value: impl Into<TValue>) {
        let value = value.into();
        unsafe {
            ptr::copy_nonoverlapping(
                (&raw const (*self.as_ptr()).key.value).cast::<u8>(),
                (&raw mut (*value.as_ptr()).value).cast::<u8>(),
                core::mem::size_of::<RawValue>(),
            );
            ptr::copy_nonoverlapping(
                (&raw const (*self.as_ptr()).key.extra).cast::<u8>(),
                (&raw mut (*value.as_ptr()).extra).cast::<u8>(),
                core::mem::size_of::<[i32; LUA_EXTRA_SIZE]>(),
            );
            (*value.as_ptr()).tt = self.key().tt();
        }
    }
}

#[allow(
    clippy::missing_safety_doc,
    reason = "LuaNodeCursor navigation is governed by Table's storage contract"
)]
impl LuaNodeCursor {
    /// Returns the current traversal address, which may be null.
    ///
    /// Table resize and rehash invalidate cursors into the old node array.
    pub const fn as_ptr(&self) -> *mut RawLuaNode {
        self.0
    }

    pub const fn from_ptr(raw: *mut RawLuaNode) -> Self {
        Self(raw)
    }

    pub const fn is_null(&self) -> bool {
        self.0.is_null()
    }

    pub unsafe fn node_unchecked(&self) -> LuaNode {
        debug_assert!(!self.is_null());
        unsafe { LuaNode::from_raw(NonNull::new_unchecked(self.0)) }
    }

    pub fn node(&self) -> Option<LuaNode> {
        NonNull::new(self.0).map(|raw| unsafe { LuaNode::from_raw(raw) })
    }

    pub unsafe fn add(self, count: usize) -> Self {
        unsafe { Self::from_ptr(self.0.add(count)) }
    }

    pub unsafe fn sub(self, count: usize) -> Self {
        unsafe { Self::from_ptr(self.0.sub(count)) }
    }

    pub unsafe fn offset(self, count: isize) -> Self {
        unsafe { Self::from_ptr(self.0.offset(count)) }
    }

    pub unsafe fn offset_from(self, other: Self) -> isize {
        unsafe { self.0.offset_from(other.0) }
    }
}
impl crate::handle::sealed::Sealed for LuaNode {}
impl RawHandle for LuaNode {
    type Raw = RawLuaNode;

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

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

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

const _: () = assert!(offset_of!(RawLuaNode, value) == 0);

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

    pub unsafe fn node(&self, index: i32) -> LuaNode {
        unsafe { self.node_cursor().add(index as usize).node_unchecked() }
    }

    pub unsafe fn node_mut(&mut self, index: i32) -> LuaNode {
        unsafe {
            LuaNode::from_raw(NonNull::new_unchecked(
                self.as_ptr()
                    .as_mut()
                    .unwrap_unchecked()
                    .node
                    .add(index as usize),
            ))
        }
    }

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

    pub unsafe fn set_metatable(&self, metatable: Option<Table>) {
        unsafe {
            self.as_ptr().as_mut().unwrap_unchecked().metatable =
                metatable.map_or(core::ptr::null_mut(), |table| table.as_ptr())
        };
    }

    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 invalidate_tm_cache(&self) {
        unsafe { self.as_ptr().as_mut().unwrap_unchecked().tm_cache = 0 };
    }

    pub unsafe fn set_node(&self, node_cursor: LuaNodeCursor) {
        unsafe { self.as_ptr().as_mut().unwrap_unchecked().node = node_cursor.as_ptr() };
    }

    pub unsafe fn set_array(&self, array_cursor: TValueCursor) {
        unsafe { self.as_ptr().as_mut().unwrap_unchecked().array = array_cursor.as_ptr() };
    }

    pub unsafe fn has_dummy_node(&self) -> bool {
        unsafe { self.node_cursor() == Self::dummy_node_cursor() }
    }

    pub unsafe fn node_index(&self, node_cursor: LuaNodeCursor) -> i32 {
        unsafe { node_cursor.offset_from(self.node_cursor()) as i32 }
    }

    /// `gval2slot`
    ///
    /// # Safety
    /// `value` must be a value pointer produced by `luaH_get` or `luaH_setslot` for this table.
    pub unsafe fn value_slot_unchecked(&self, value: TValue) -> i32 {
        unsafe {
            let value_addr = value.as_ptr() as usize;
            let node_addr = self.as_ptr().as_ref().unwrap_unchecked().node as usize;
            (value_addr.wrapping_sub(node_addr) / mem::size_of::<RawLuaNode>()) as i32
        }
    }

    pub unsafe fn node_cursor(&self) -> LuaNodeCursor {
        unsafe { LuaNodeCursor::from_ptr(self.as_ptr().as_ref().unwrap_unchecked().node) }
    }

    pub unsafe fn array_cursor(&self) -> TValueCursor {
        unsafe { TValueCursor::from_ptr(self.as_ptr().as_ref().unwrap_unchecked().array) }
    }

    pub unsafe fn array_slot(&self, index: usize) -> TValue {
        unsafe { self.array_cursor().add(index).value_unchecked() }
    }

    pub unsafe fn array_slot_for_key(&self, key: i32) -> Option<TValue> {
        unsafe {
            ((key as u32).wrapping_sub(1)
                < self.as_ptr().as_ref().unwrap_unchecked().size_array as u32)
                .then(|| self.array_slot((key - 1) as usize))
        }
    }

    pub unsafe fn node_count(&self) -> usize {
        unsafe { 1usize << self.as_ptr().as_ref().unwrap_unchecked().lsize_node }
    }

    pub unsafe fn hash_mask(&self) -> usize {
        unsafe { self.node_count() - 1 }
    }

    pub unsafe fn node_mask_8(&self) -> usize {
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().node_mask_8 as usize }
    }

    pub unsafe fn array_storage_size(&self) -> usize {
        unsafe {
            self.as_ptr().as_ref().unwrap_unchecked().size_array as usize
                * mem::size_of::<RawTValue>()
        }
    }

    pub unsafe fn node_storage_size(&self) -> usize {
        unsafe {
            if self.has_dummy_node() {
                0
            } else {
                self.node_count() * mem::size_of::<RawLuaNode>()
            }
        }
    }

    pub unsafe fn allocation_size(&self) -> usize {
        unsafe {
            mem::size_of::<RawLuaTable>() + self.array_storage_size() + self.node_storage_size()
        }
    }

    pub unsafe fn gc_work_size(&self, count_dummy_node: bool) -> usize {
        unsafe {
            let node_size = if count_dummy_node {
                self.node_count() * mem::size_of::<RawLuaNode>()
            } else {
                self.node_storage_size()
            };
            mem::size_of::<RawLuaTable>() + self.array_storage_size() + node_size
        }
    }

    pub unsafe fn init_empty_storage(&self) {
        unsafe {
            let table_ref = self.as_ptr().as_mut().unwrap_unchecked();
            table_ref.array = core::ptr::null_mut();
            table_ref.size_array = 0;
            table_ref.lsize_node = 0;
            table_ref.readonly = 0;
            table_ref.safe_env = 0;
            table_ref.node_mask_8 = 0;
            table_ref.node = Self::dummy_node_ptr();
            table_ref.gc_list = core::ptr::null_mut();
            table_ref.free = RawLuaTableFree { last_free: 0 };
        }
    }

    pub unsafe fn free_storage(&self, thread: &Thread) {
        unsafe {
            let table_ref = self.as_ptr().as_ref().unwrap_unchecked();
            if !self.has_dummy_node() {
                thread.free_array(
                    self.node_cursor().as_ptr(),
                    self.node_count(),
                    table_ref.memcat,
                );
            }

            if !table_ref.array.is_null() {
                thread.free_array(
                    table_ref.array,
                    table_ref.size_array as usize,
                    table_ref.memcat,
                );
            }
        }
    }

    pub unsafe fn maybe_set_aboundary(&self, value: i32) {
        if unsafe { self.as_ptr().as_ref().unwrap_unchecked().free.aboundary } <= 0 {
            unsafe {
                self.as_ptr().as_mut().unwrap_unchecked().free =
                    RawLuaTableFree { aboundary: -value }
            };
        }
    }

    /// `updateaboundary`
    pub(super) fn update_aboundary(&self, boundary: i32) -> i32 {
        let size_array = unsafe { self.as_ptr().as_ref().unwrap_unchecked().size_array };
        let array = unsafe { self.array_cursor() };
        let boundary_slot_is_nil = unsafe { array.add((boundary - 1) as usize).is_nil_unchecked() };
        if boundary < size_array && boundary_slot_is_nil {
            let previous_slot_is_set = if boundary >= 2 {
                unsafe { !array.add((boundary - 2) as usize).is_nil_unchecked() }
            } else {
                false
            };
            if previous_slot_is_set {
                unsafe {
                    self.maybe_set_aboundary(boundary - 1);
                }
                return boundary - 1;
            }
        } else {
            let next_slot_is_set = if boundary + 1 < size_array {
                unsafe { !array.add(boundary as usize).is_nil_unchecked() }
            } else {
                false
            };
            let slot_after_next_is_nil = if boundary + 1 < size_array {
                unsafe { array.add((boundary + 1) as usize).is_nil_unchecked() }
            } else {
                false
            };
            if next_slot_is_set && slot_after_next_is_nil {
                unsafe {
                    self.maybe_set_aboundary(boundary + 1);
                }
                return boundary + 1;
            }
        }

        0
    }

    pub unsafe fn get_aboundary(&self) -> i32 {
        let aboundary = unsafe { self.as_ptr().as_ref().unwrap_unchecked().free.aboundary };
        if aboundary < 0 {
            -aboundary
        } else {
            unsafe { self.as_ptr().as_ref().unwrap_unchecked().size_array }
        }
    }
}