luau-vm 0.732.0

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

use crate::gc::{GcObject, RawGcObject};
use crate::handle::RawHandle;
use crate::string::{RawTString, TString};
use crate::value::{RawTValue, TValue, TValueCursor};

#[repr(C)]
#[derive(Clone, Copy)]
pub struct FeedbackVectorSlotCallTarget {
    pub pc: u32,
    pub proto: u32,
    pub hits: u32,
}

#[repr(C)]
#[derive(Clone, Copy)]
pub union FeedbackVectorSlotData {
    pub call_target: FeedbackVectorSlotCallTarget,
}

#[repr(C)]
pub struct FeedbackVectorSlot {
    pub kind: i32,
    pub data: FeedbackVectorSlotData,
}

#[repr(C)]
pub struct RawProto {
    pub tt: u8,
    pub marked: u8,
    pub memcat: u8,
    pub n_ups: u8,
    pub num_params: u8,
    pub is_vararg: u8,
    pub max_stack_size: u8,
    pub flags: u8,
    pub k: *mut RawTValue,
    pub code: *mut u32,
    pub p: *mut *mut RawProto,
    pub code_entry: *const u32,
    pub exec_data: *mut (),
    pub exec_target: usize,
    pub line_info: *mut u8,
    pub abs_line_info: *mut i32,
    pub loc_vars: *mut RawLocVar,
    pub upvalues: *mut *mut RawTString,
    pub source: *mut RawTString,
    pub debug_name: *mut RawTString,
    pub debug_insn: *mut u8,
    pub type_info: *mut u8,
    pub userdata: *mut (),
    pub gc_list: *mut RawGcObject,
    pub size_code: i32,
    pub size_p: i32,
    pub size_loc_vars: i32,
    pub size_upvalues: i32,
    pub size_k: i32,
    pub size_line_info: i32,
    pub line_gap_log2: i32,
    pub line_defined: i32,
    pub bytecode_id: i32,
    pub size_type_info: i32,
    pub feedback_vec: *mut FeedbackVectorSlot,
    pub feedback_vec_size: u32,
    pub fun_id: u32,
    pub cost: u64,
}

#[repr(C)]
pub struct RawLocVar {
    pub var_name: *mut RawTString,
    pub start_pc: i32,
    pub end_pc: i32,
    pub reg: u8,
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
/// Non-owning identity of a VM prototype record.
///
/// # Safety model for unsafe methods
///
/// The proto and all referenced arrays and objects must remain live. Program
/// counters and indices must be in bounds for their corresponding arrays, and
/// code mutation must preserve a valid bytecode graph.
pub struct Proto {
    raw: NonNull<RawProto>,
}

#[derive(Clone, Copy, PartialEq, Eq)]
#[repr(transparent)]
/// Non-owning view of a prototype local-variable record.
///
/// Unsafe operations require the owning proto and local-variable array to
/// remain live for the entire use of this copied handle.
pub struct LocVar {
    raw: NonNull<RawLocVar>,
}

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

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

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

    pub unsafe fn upvalue_name(&self, index: usize) -> Option<TString> {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        if index >= proto.size_upvalues as usize {
            None
        } else {
            NonNull::new(unsafe { *proto.upvalues.add(index) })
                .map(|raw| unsafe { TString::from_raw(raw) })
        }
    }

    pub unsafe fn child_proto(&self, index: usize) -> Option<Self> {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        if index >= proto.size_p as usize {
            None
        } else {
            NonNull::new(unsafe { *proto.p.add(index) }).map(|raw| unsafe { Self::from_raw(raw) })
        }
    }

    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 loc_var(&self, index: usize) -> Option<LocVar> {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        if index >= proto.size_loc_vars as usize {
            None
        } else {
            Some(unsafe { LocVar::from_raw(NonNull::new_unchecked(proto.loc_vars.add(index))) })
        }
    }

    /// `sizeproto`
    pub unsafe fn size(&self) -> usize {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };

        size_of::<RawProto>()
            + size_of::<u32>() * proto.size_code as usize
            + size_of::<*mut RawProto>() * proto.size_p as usize
            + size_of::<RawTValue>() * proto.size_k as usize
            + proto.size_line_info as usize
            + size_of::<RawLocVar>() * proto.size_loc_vars as usize
            + size_of::<*mut RawTString>() * proto.size_upvalues as usize
            + proto.size_type_info as usize
    }

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

    pub unsafe fn constant(&self, index: usize) -> TValue {
        debug_assert!(index < unsafe { self.as_ptr().as_ref().unwrap_unchecked().size_k as usize });
        unsafe { self.constants().add(index).value_unchecked() }
    }

    /// `luaG_getline`
    pub unsafe fn get_line(&self, pc: i32) -> i32 {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        debug_assert!(pc >= 0);
        debug_assert!(pc < proto.size_code);

        unsafe {
            if proto.line_info.is_null() {
                0
            } else {
                *proto.abs_line_info.add(pc as usize >> proto.line_gap_log2)
                    + i32::from(*proto.line_info.add(pc as usize))
            }
        }
    }

    /// `pcRel`
    pub unsafe fn pc_rel(&self, saved_pc: *const u32) -> i32 {
        let code = unsafe { self.as_ptr().as_ref().unwrap_unchecked().code };
        if saved_pc.is_null() || saved_pc == code {
            0
        } else {
            unsafe { saved_pc.offset_from(code) as i32 - 1 }
        }
    }

    /// `luaF_getlocal`
    pub unsafe fn get_local(&self, mut local_number: i32, pc: i32) -> Option<LocVar> {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };

        for index in 0..proto.size_loc_vars as usize {
            let local = unsafe { self.loc_var(index).unwrap_unchecked() };
            let local_ref = unsafe { local.as_ptr().as_ref().unwrap_unchecked() };
            if pc >= local_ref.start_pc && pc < local_ref.end_pc {
                local_number -= 1;
                if local_number == 0 {
                    return Some(local);
                }
            }
        }

        None
    }

    /// `luaF_findlocal`
    pub unsafe fn find_local(&self, local_reg: i32, pc: i32) -> Option<LocVar> {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };

        for index in 0..proto.size_loc_vars as usize {
            let local = unsafe { self.loc_var(index).unwrap_unchecked() };
            let local_ref = unsafe { local.as_ptr().as_ref().unwrap_unchecked() };
            if i32::from(local_ref.reg) == local_reg
                && pc >= local_ref.start_pc
                && pc < local_ref.end_pc
            {
                return Some(local);
            }
        }

        None
    }

    pub unsafe fn feedback_slot(&self, index: usize) -> Option<NonNull<FeedbackVectorSlot>> {
        let proto = unsafe { self.as_ptr().as_ref().unwrap_unchecked() };
        if index >= proto.feedback_vec_size as usize {
            None
        } else {
            NonNull::new(unsafe { proto.feedback_vec.add(index) })
        }
    }

    pub unsafe fn code_at(&self, index: usize) -> u32 {
        unsafe { *self.as_ptr().as_ref().unwrap_unchecked().code.add(index) }
    }

    pub unsafe fn set_code_at(&self, index: usize, value: u32) {
        unsafe {
            *self.as_ptr().as_ref().unwrap_unchecked().code.add(index) = value;
        }
    }
}

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

    pub unsafe fn name(&self) -> Option<TString> {
        unsafe {
            NonNull::new(self.as_ptr().as_ref().unwrap_unchecked().var_name)
                .map(|raw| TString::from_raw(raw))
        }
    }
}
impl crate::handle::sealed::Sealed for Proto {}
impl crate::handle::sealed::Sealed for LocVar {}
impl RawHandle for Proto {
    type Raw = RawProto;

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

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

impl RawHandle for LocVar {
    type Raw = RawLocVar;

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

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