luau-vm 0.732.0

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

use luau_common::{BStr, ByteSlice};

use crate::VmErrorResult;
use crate::gc::{FIXED_BIT, GcObject, bit_mask};
use crate::handle::RawHandle;
use crate::handle::sealed::Sealed;
use crate::memory::{LuaPage, MemoryRuntime};
use crate::thread::Thread;
use crate::types::LUA_TSTRING;

pub const MAX_STRING_SIZE: usize = 1 << 30;
pub const ATOM_UNDEFINED: i16 = i16::MIN;

/// `LUA_MINSTRTABSIZE`
pub(crate) const LUA_MIN_STRING_TABLE_SIZE: usize = 32;

/// `luaS_hash`
pub fn hash(bytes: &[u8]) -> u32 {
    let mut a = 0u32;
    let mut b = 0u32;
    let mut hash = bytes.len() as u32;
    let mut cursor = bytes;
    let mut len = bytes.len();

    while len >= 32 {
        let block0 = u32::from_ne_bytes(cursor[..4].try_into().unwrap());
        let block1 = u32::from_ne_bytes(cursor[4..8].try_into().unwrap());
        let block2 = u32::from_ne_bytes(cursor[8..12].try_into().unwrap());

        a = a.wrapping_add(block0);
        b = b.wrapping_add(block1);
        hash = hash.wrapping_add(block2);

        a ^= hash;
        a = a.wrapping_sub(hash.rotate_right(14));
        b ^= a;
        b = b.wrapping_sub(a.rotate_right(11));
        hash ^= b;
        hash = hash.wrapping_sub(b.rotate_right(25));

        cursor = &cursor[12..];
        len -= 12;
    }

    for index in (0..len).rev() {
        let byte = cursor[index];
        hash ^= hash
            .wrapping_shl(5)
            .wrapping_add(hash.wrapping_shr(2))
            .wrapping_add(u32::from(byte));
    }

    hash
}

#[repr(C)]
pub struct StringTable {
    pub hash: *mut *mut RawTString,
    pub n_use: u32,
    pub size: i32,
}

#[repr(C)]
pub struct RawTString {
    pub tt: u8,
    pub marked: u8,
    pub memcat: u8,
    pub atom: i16,
    pub next: *mut RawTString,
    pub hash: u32,
    pub len: u32,
    pub data: [u8; 0],
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
pub struct TString {
    raw: NonNull<RawTString>,
}

/// Unstable interned-string allocation capability.
///
/// # Safety
///
/// The thread and all string/page handles must be live and belong to the same
/// VM. Sizes and pages must match their allocations, and callers must preserve
/// interning-table, atom, rooting, and GC invariants.
#[allow(
    clippy::missing_safety_doc,
    reason = "all methods share the capability-level safety contract"
)]
pub trait StringRuntime: Sealed {
    /// `luaS_resize`
    unsafe fn resize(&self, new_size: i32) -> VmErrorResult;

    /// `luaS_newlstr`
    unsafe fn intern_string(&self, bytes: &BStr) -> VmErrorResult<TString>;

    /// `luaS_free`
    unsafe fn free_string(&self, string: TString, page: LuaPage);

    /// `luaS_updateatom`
    unsafe fn update_atom(&self, string: TString);

    /// `luaS_bufstart`
    unsafe fn buffer_start(&self, size: usize) -> VmErrorResult<TString>;

    /// `luaS_buffinish`
    unsafe fn buffer_finish(&self, string: TString) -> VmErrorResult<TString>;
}

impl TString {
    /// `luaS_fix`
    ///
    /// # Safety
    ///
    /// The string must still belong to a live VM, and the caller must ensure
    /// that no concurrent or reentrant collector operation accesses its mark
    /// bits while they are updated.
    pub unsafe fn fix(&self) {
        unsafe {
            self.as_ptr().as_mut().unwrap_unchecked().marked |= bit_mask(FIXED_BIT);
        }
    }

    /// Constructs a non-owning string handle from a raw VM record.
    ///
    /// # Safety
    ///
    /// `raw` must address a live `RawTString` owned by the VM in which the
    /// handle will be used. The handle does not root the string or extend its
    /// lifetime.
    pub const unsafe fn from_raw(raw: NonNull<RawTString>) -> Self {
        Self { raw }
    }

    /// Constructs a non-owning string handle from a VM string record.
    ///
    /// # Safety
    ///
    /// `raw` must be a live, VM-owned string record. The returned handle may
    /// not outlive that record and does not keep it reachable from the GC.
    pub unsafe fn from_ref(raw: &RawTString) -> Self {
        Self {
            raw: NonNull::from(raw),
        }
    }

    pub fn atom(&self) -> i16 {
        unsafe { (*self.as_ptr()).atom }
    }

    pub fn set_atom(&self, atom: i16) {
        unsafe {
            (*self.as_ptr()).atom = atom;
        }
    }

    pub const fn size_string(len: usize) -> usize {
        core::mem::offset_of!(RawTString, data) + len
    }

    pub fn data_ptr(&self) -> *const u8 {
        unsafe { (&raw const (*self.as_ptr()).data).cast::<u8>() }
    }

    /// Returns the address of the string's inline byte storage.
    ///
    /// # Safety
    ///
    /// The string must be live. Mutating interned string bytes can invalidate
    /// hashes and table invariants, so writes are only valid during the VM's
    /// owned string-construction protocol before the string is published.
    pub unsafe fn data_mut_ptr(&self) -> *mut u8 {
        unsafe { (&raw mut (*self.as_ptr()).data).cast::<u8>() }
    }

    /// Borrows the string's bytes for the duration of this handle borrow.
    ///
    /// # Safety
    ///
    /// The string must remain live and its storage must not be freed for the
    /// returned borrow. Because this handle does not root the string, the
    /// caller must also prevent VM transitions that could collect it.
    pub unsafe fn as_bytes(&self) -> &[u8] {
        unsafe { core::slice::from_raw_parts(self.data_ptr(), (*self.as_ptr()).len as usize) }
    }

    /// Borrows the string's bytes as a `BStr`.
    ///
    /// # Safety
    ///
    /// The string must remain live and its storage must not be freed for the
    /// returned borrow. Because this handle does not root the string, the
    /// caller must also prevent VM transitions that could collect it.
    pub unsafe fn as_bstr(&self) -> &BStr {
        unsafe { self.as_bytes().as_bstr() }
    }
}

impl crate::handle::sealed::Sealed for TString {}

impl RawHandle for TString {
    type Raw = RawTString;

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

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

fn bucket_index(hash: u32, size: i32) -> usize {
    debug_assert!(size > 0);
    debug_assert!(size & (size - 1) == 0);
    (hash & (size as u32 - 1)) as usize
}

impl StringRuntime for Thread {
    /// `luaS_resize`
    unsafe fn resize(&self, new_size: i32) -> VmErrorResult {
        unsafe {
            let new_hash = self.new_array::<*mut RawTString>(new_size as usize, 0)?;

            for index in 0..new_size as usize {
                new_hash.add(index).write(ptr::null_mut());
            }

            let global = self.global();
            let string_table = &global.as_ptr().as_ref().unwrap_unchecked().string_table;
            let old_hash = string_table.hash;
            let old_size = string_table.size;

            for index in 0..old_size as usize {
                let mut string = NonNull::new(old_hash.add(index).read());
                while let Some(current) = string {
                    let next = NonNull::new(current.as_ref().next);
                    let bucket = bucket_index(current.as_ref().hash, new_size);

                    current.as_ptr().as_mut().unwrap_unchecked().next = new_hash.add(bucket).read();
                    new_hash.add(bucket).write(current.as_ptr());

                    string = next;
                }
            }

            if !old_hash.is_null() {
                self.free_array(old_hash, old_size as usize, 0);
            }

            let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
            string_table.size = new_size;
            string_table.hash = new_hash;
        }
        Ok(())
    }

    /// `luaS_newlstr`
    unsafe fn intern_string(&self, bytes: &BStr) -> VmErrorResult<TString> {
        let bytes = bytes.as_bytes();
        let len = bytes.len();
        let hash = crate::string::hash(bytes);
        unsafe {
            let global = self.global();
            {
                let string_table = &global.as_ptr().as_ref().unwrap_unchecked().string_table;
                debug_assert!(string_table.size > 0);

                let bucket = bucket_index(hash, string_table.size);
                let mut entry = NonNull::new(string_table.hash.add(bucket).read());

                while let Some(current) = entry {
                    let current_string = TString::from_raw(current);
                    if current.as_ref().len as usize == len && current_string.as_bytes() == bytes {
                        let mut object: GcObject = current_string.into();
                        if global.is_dead(object) {
                            object.change_white();
                        }

                        return Ok(current_string);
                    }

                    entry = NonNull::new(current.as_ref().next);
                }
            }

            if len > MAX_STRING_SIZE {
                return self.too_big();
            }

            let active_memcat = self.as_ptr().as_ref().unwrap_unchecked().active_memcat;
            let string_handle =
                self.new_gco::<TString>(TString::size_string(len), active_memcat)?;
            GcObject::from(string_handle).init_header(self, LUA_TSTRING as u8);
            let string_ref = string_handle.as_ptr().as_mut().unwrap_unchecked();
            string_ref.atom = ATOM_UNDEFINED;
            string_ref.hash = hash;
            string_ref.len = len as u32;
            ptr::copy_nonoverlapping(bytes.as_ptr(), string_handle.data_mut_ptr(), len);
            let (should_resize, next_size) = {
                let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
                let bucket = bucket_index(hash, string_table.size);
                string_ref.next = string_table.hash.add(bucket).read();
                string_table.hash.add(bucket).write(string_handle.as_ptr());

                string_table.n_use += 1;
                (
                    string_table.n_use > string_table.size as u32
                        && string_table.size <= i32::MAX / 2,
                    string_table.size * 2,
                )
            };

            if should_resize {
                self.resize(next_size)?;
            }

            Ok(string_handle)
        }
    }

    /// `luaS_free`
    unsafe fn free_string(&self, string: TString, page: LuaPage) {
        unsafe {
            let len = string.as_ptr().as_ref().unwrap_unchecked().len as usize;
            let memcat = string.as_ptr().as_ref().unwrap_unchecked().memcat;

            let global = self.global();
            let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
            let bucket = bucket_index(
                string.as_ptr().as_ref().unwrap_unchecked().hash,
                string_table.size,
            );
            let mut slot = string_table.hash.add(bucket);
            let mut found = false;
            let string_raw = string.as_ptr();

            while let Some(current) = NonNull::new(slot.read()) {
                if current.as_ptr() == string_raw {
                    slot.write(current.as_ref().next);
                    found = true;
                    string_table.n_use -= 1;
                    break;
                }

                slot = &raw mut (*current.as_ptr()).next;
            }

            debug_assert!(found || string.as_ptr().as_ref().unwrap_unchecked().next.is_null());

            self.free_gco(string.into(), TString::size_string(len), memcat, page);
        }
    }

    /// `luaS_updateatom`
    unsafe fn update_atom(&self, string: TString) {
        unsafe {
            if string.atom() != ATOM_UNDEFINED {
                return;
            }

            let global = self.global();
            let atom = if let Some(user_atom) = global.user_atom_callback() {
                user_atom(self, string.as_bstr())
            } else {
                -1
            };
            string.set_atom(atom);
        }
    }

    /// `luaS_bufstart`
    unsafe fn buffer_start(&self, size: usize) -> VmErrorResult<TString> {
        if size > MAX_STRING_SIZE {
            return unsafe { self.too_big() };
        }

        unsafe {
            let string = self.new_gco::<TString>(
                TString::size_string(size),
                self.as_ptr().as_ref().unwrap_unchecked().active_memcat,
            )?;
            GcObject::from(string).init_header(self, LUA_TSTRING as u8);
            let string_ref = string.as_ptr().as_mut().unwrap_unchecked();
            string_ref.atom = ATOM_UNDEFINED;
            string_ref.hash = 0;
            string_ref.len = size as u32;
            string_ref.next = ptr::null_mut();

            Ok(string)
        }
    }

    /// `luaS_buffinish`
    unsafe fn buffer_finish(&self, string: TString) -> VmErrorResult<TString> {
        let bytes = unsafe { string.as_bytes() };
        let hash = crate::string::hash(bytes);
        unsafe {
            let global = self.global();
            let string_table = &mut global.as_ptr().as_mut().unwrap_unchecked().string_table;
            let bucket = bucket_index(hash, string_table.size);

            let mut entry = NonNull::new(string_table.hash.add(bucket).read());
            while let Some(current) = entry {
                let current_string = TString::from_raw(current);
                if current_string.as_ptr().as_ref().unwrap_unchecked().len
                    == string.as_ptr().as_ref().unwrap_unchecked().len
                    && current_string.as_bytes() == bytes
                {
                    let mut object: GcObject = current_string.into();
                    if global.is_dead(object) {
                        object.change_white();
                    }

                    return Ok(current_string);
                }

                entry = NonNull::new(current.as_ref().next);
            }

            let string_ref = string.as_ptr().as_mut().unwrap_unchecked();
            string_ref.hash = hash;
            string_ref.atom = ATOM_UNDEFINED;
            string_ref.next = string_table.hash.add(bucket).read();
            string_table.hash.add(bucket).write(string.as_ptr());

            string_table.n_use += 1;
            if string_table.n_use > string_table.size as u32 && string_table.size <= i32::MAX / 2 {
                self.resize(string_table.size * 2)?;
            }

            Ok(string)
        }
    }
}