luau-vm 0.732.0

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

use crate::gc::{GcObject, RawGcObject};
use crate::handle::RawHandle;
use crate::state::{GlobalState, VmAllocator};
use crate::types::LUA_TNIL;

use super::{GCO_LINK_OFFSET, LUA_PAGE_PADDING};

/// Mutable traversal state for the allocated objects in one VM page.
///
/// Unsafe traversal requires the page to remain live and unmodified, with a
/// valid block size and end pointer, until the walk completes.
pub struct GcoPageWalk {
    pub(super) current: *mut u8,
    pub(super) end: *mut u8,
    pub(super) block_size: i32,
    pub(super) steps: i32,
}
#[allow(
    clippy::missing_safety_doc,
    reason = "GcoPageWalk's shared traversal contract is documented on GcoPageWalk"
)]
impl GcoPageWalk {
    pub unsafe fn next(&mut self) -> Option<GcObject> {
        while self.current != self.end {
            let block = self.current.cast::<RawGcObject>();
            self.current = unsafe { self.current.add(self.block_size as usize) };
            self.steps += 1;

            if unsafe { (*block).tt } == LUA_TNIL as u8 {
                continue;
            }

            return Some(unsafe { GcObject::from_raw(NonNull::new_unchecked(block)) });
        }

        None
    }

    pub fn steps(&self) -> i32 {
        self.steps
    }
}

#[repr(C)]
pub struct RawLuaPage {
    pub prev: *mut RawLuaPage,
    pub next: *mut RawLuaPage,
    pub list_prev: *mut RawLuaPage,
    pub list_next: *mut RawLuaPage,
    pub page_size: i32,
    pub block_size: i32,
    pub free_list: *mut u8,
    pub free_next: i32,
    pub busy_blocks: i32,
    pub padding: [u8; LUA_PAGE_PADDING],
    pub data: [u8; 1],
}

#[derive(Clone, Copy)]
#[repr(transparent)]
pub(crate) struct LuaPageList {
    head: NonNull<*mut RawLuaPage>,
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
/// Non-owning identity of a VM allocation page.
///
/// # Safety model for unsafe methods
///
/// The page and its allocation must remain live. Block sizes, offsets, and
/// traversal metadata must match the page class; page-list or allocator
/// mutation can invalidate the handle and every pointer derived from it.
pub struct LuaPage {
    raw: NonNull<RawLuaPage>,
}

impl LuaPageList {
    pub(crate) const unsafe fn from_raw(head: NonNull<*mut RawLuaPage>) -> Self {
        Self { head }
    }

    pub(crate) unsafe fn get(&self) -> *mut RawLuaPage {
        unsafe { *self.head.as_ptr() }
    }

    pub(crate) unsafe fn set(&self, page: *mut RawLuaPage) {
        unsafe {
            *self.head.as_ptr() = page;
        }
    }

    pub(crate) unsafe fn link_page(&self, page: LuaPage) {
        unsafe {
            let page_ref = page.as_ptr().as_mut().unwrap_unchecked();
            page_ref.list_next = self.get();
            if let Some(mut next) = NonNull::new(page_ref.list_next) {
                next.as_mut().list_prev = page.as_ptr();
            }
            self.set(page.as_ptr());
        }
    }

    pub(crate) unsafe fn unlink_page(&self, page: LuaPage) {
        unsafe {
            let page_ref = page.as_ptr().as_mut().unwrap_unchecked();
            if let Some(mut list_next) = NonNull::new(page_ref.list_next) {
                list_next.as_mut().list_prev = page_ref.list_prev;
            }

            if let Some(mut list_prev) = NonNull::new(page_ref.list_prev) {
                list_prev.as_mut().list_next = page_ref.list_next;
            } else if self.get() == page.as_ptr() {
                self.set(page_ref.list_next);
            }

            page_ref.list_prev = ptr::null_mut();
            page_ref.list_next = ptr::null_mut();
        }
    }
}

impl GlobalState {
    pub(crate) fn allocator(&self) -> NonNull<VmAllocator> {
        unsafe { self.as_ptr().as_ref().unwrap_unchecked().allocator }
    }

    pub(crate) unsafe fn free_page_list(&self, size_class_index: usize) -> LuaPageList {
        unsafe {
            let slots = (&raw mut (*self.as_ptr()).free_pages).cast::<*mut RawLuaPage>();
            LuaPageList::from_raw(NonNull::new_unchecked(slots.add(size_class_index)))
        }
    }

    pub(crate) unsafe fn free_gco_page_list(&self, size_class_index: usize) -> LuaPageList {
        unsafe {
            let slots = (&raw mut (*self.as_ptr()).free_gco_pages).cast::<*mut RawLuaPage>();
            LuaPageList::from_raw(NonNull::new_unchecked(slots.add(size_class_index)))
        }
    }

    pub(crate) unsafe fn all_page_list(&self) -> LuaPageList {
        unsafe {
            LuaPageList::from_raw(NonNull::new_unchecked(&raw mut (*self.as_ptr()).all_pages))
        }
    }

    pub(crate) unsafe fn all_gco_page_list(&self) -> LuaPageList {
        unsafe {
            LuaPageList::from_raw(NonNull::new_unchecked(
                &raw mut (*self.as_ptr()).all_gco_pages,
            ))
        }
    }

    pub fn all_gco_pages(&self) -> Option<LuaPage> {
        unsafe { NonNull::new((*self.as_ptr()).all_gco_pages).map(|raw| LuaPage::from_raw(raw)) }
    }

    pub fn sweep_gco_page(&self) -> Option<LuaPage> {
        unsafe { NonNull::new((*self.as_ptr()).sweep_gco_page).map(|raw| LuaPage::from_raw(raw)) }
    }

    pub fn set_sweep_gco_page(&self, page: Option<LuaPage>) {
        unsafe {
            (*self.as_ptr()).sweep_gco_page = page.map_or(ptr::null_mut(), |page| page.as_ptr());
        }
    }
}

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

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

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

    /// Returns a block address at `offset` bytes from the page data.
    ///
    /// # Safety
    ///
    /// `offset` must be non-negative, aligned for the page's block layout, and
    /// identify a complete block within this page's allocation.
    pub unsafe fn block(&self, offset: i32) -> *mut u8 {
        unsafe { self.data_ptr_mut().add(offset as usize) }
    }

    pub fn page_blocks(&self) -> i32 {
        let raw = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        ((raw.page_size as usize - offset_of!(RawLuaPage, data)) / raw.block_size as usize) as i32
    }

    /// # Safety
    ///
    /// `block` must address a live allocator block whose leading storage is a
    /// pointer-sized metadata slot.
    pub unsafe fn metadata_slot(block: *mut u8) -> *mut *mut u8 {
        block.cast::<*mut u8>()
    }

    /// # Safety
    ///
    /// `block` must address a live GC block large enough to contain the
    /// collector free-list link at `GCO_LINK_OFFSET`.
    pub unsafe fn free_gco_link_slot(block: *mut u8) -> *mut *mut u8 {
        unsafe { block.add(GCO_LINK_OFFSET).cast::<*mut u8>() }
    }

    pub unsafe fn gco_walk(&self) -> (GcoPageWalk, i32) {
        let page = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        let block_count =
            (page.page_size as usize - offset_of!(RawLuaPage, data)) / page.block_size as usize;

        (
            GcoPageWalk {
                current: unsafe {
                    self.data_ptr_mut()
                        .add((page.free_next + page.block_size) as usize)
                },
                end: unsafe {
                    self.data_ptr_mut()
                        .add(block_count * page.block_size as usize)
                },
                block_size: page.block_size,
                steps: 0,
            },
            page.busy_blocks,
        )
    }

    /// `luaM_getnextpage`
    pub unsafe fn next_page(&self) -> Option<LuaPage> {
        NonNull::new(unsafe { self.as_ptr().as_ref().unwrap_unchecked().list_next })
            .map(|page| unsafe { Self::from_raw(page) })
    }

    /// `luaM_visitpage`
    pub unsafe fn visit_page(
        &self,
        context: *mut (),
        visitor: unsafe fn(*mut (), LuaPage, GcObject) -> bool,
    ) {
        unsafe {
            let (mut walk, mut busy) = self.gco_walk();
            while let Some(gco) = walk.next() {
                let should_delete = visitor(context, *self, gco);
                if should_delete {
                    busy -= 1;
                    if busy == 0 {
                        break;
                    }
                }
            }
        }
    }
}

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

impl RawHandle for LuaPage {
    type Raw = RawLuaPage;

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

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