luau-vm 0.732.0

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

use luau_bytecode::model::Instruction;
use luau_bytecode::opcodes::Opcode;
use luau_common::{BStr, BString, ByteSlice};

use super::stack::RawStackAccess;
use super::{LuaStringBuilder, LuaStringBuilderStorage, Thread};
use crate::VmErrorResult;
use crate::debug::{
    DebugRuntime, LuaCounterFunction, LuaCounterValue, LuaCoverage, LuaDebug, chunk_id, current_pc,
    lua_proto,
};
use crate::function::{Closure, Proto};
use crate::gc::GcBarrier;
use crate::handle::RawHandle;
use crate::memory::MemoryRuntime;
use crate::state::{CallInfo, RawCallInfo, ThreadState};
use crate::string::LuaString;
use crate::value::TValue;
use crate::vm::VmCallFrame;

// Debug inspection
fn get_func_name(closure: Closure) -> Option<LuaString> {
    if unsafe { closure.is_native() } {
        unsafe { closure.native_debug_name() }
    } else {
        let proto = unsafe { closure.proto().unwrap_unchecked() };
        unsafe { proto.debug_name().map(LuaString::from_interned) }
    }
}

/// `auxgetinfo`
fn aux_get_info(
    what: &[u8],
    ar: &mut LuaDebug,
    closure: Closure,
    call_info: Option<*mut RawCallInfo>,
) -> bool {
    let mut push_function = false;

    for byte in what {
        match *byte {
            b's' => {
                if unsafe { closure.is_native() } {
                    ar.source = LuaString::from_static(b"=[C]".as_bstr());
                    ar.what = LuaString::from_static(b"C".as_bstr());
                    ar.linedefined = -1;
                    ar.set_short_src(b"[C]");
                } else {
                    let proto = unsafe { closure.proto().unwrap_unchecked() };
                    let source = unsafe { proto.source().unwrap_unchecked() };
                    ar.source = LuaString::from_interned(source);
                    ar.what = LuaString::from_static(b"Lua".as_bstr());
                    ar.linedefined =
                        unsafe { proto.as_ptr().as_ref().unwrap_unchecked().line_defined };

                    let mut short_src = [0u8; crate::debug::LUA_ID_SIZE];
                    let short_src = chunk_id(&mut short_src, unsafe { source.as_bytes() });
                    ar.set_short_src(short_src);
                }
            }
            b'l' => {
                ar.currentline = if let Some(call_info) = call_info {
                    let call_info =
                        unsafe { CallInfo::from_raw(NonNull::new_unchecked(call_info)) };
                    if let Some(proto) = lua_proto(call_info) {
                        unsafe { proto.get_line(current_pc(call_info, proto)) }
                    } else {
                        -1
                    }
                } else if unsafe { closure.is_native() } {
                    -1
                } else {
                    unsafe {
                        closure
                            .proto()
                            .unwrap_unchecked()
                            .as_ptr()
                            .as_ref()
                            .unwrap_unchecked()
                            .line_defined
                    }
                };
            }
            b'u' => {
                ar.nupvals = unsafe { closure.as_ptr().as_ref().unwrap_unchecked().n_upvalues };
            }
            b'a' => {
                if unsafe { closure.is_native() } {
                    ar.is_vararg = true;
                    ar.nparams = 0;
                } else {
                    let proto = unsafe { closure.proto().unwrap_unchecked() };
                    let proto = unsafe { proto.as_ptr().as_ref().unwrap_unchecked() };
                    ar.is_vararg = proto.is_vararg != 0;
                    ar.nparams = proto.num_params;
                }
            }
            b'n' => {
                ar.name = get_func_name(closure);
            }
            b'f' => {
                push_function = true;
            }
            _ => {}
        }
    }

    push_function
}

/// `aux_upvalue`
fn aux_upvalue(function: TValue, upvalue_index: i32) -> Option<(LuaString, TValue)> {
    if !function.is_function() {
        return None;
    }

    let closure = function.closure_value();
    if unsafe { closure.is_native() } {
        if !(1..=i32::from(unsafe { closure.as_ptr().as_ref().unwrap_unchecked().n_upvalues }))
            .contains(&upvalue_index)
        {
            return None;
        }

        let value = unsafe { closure.native_upvalue((upvalue_index - 1) as usize) };
        return Some((LuaString::from_static(b"".as_bstr()), value));
    }

    let proto = unsafe { closure.proto().unwrap_unchecked() };
    if !(1..=i32::from(unsafe { proto.as_ptr().as_ref().unwrap_unchecked().n_ups }))
        .contains(&upvalue_index)
    {
        return None;
    }

    let upref = unsafe { closure.lua_upvalue_ref((upvalue_index - 1) as usize) };
    let value = if upref.is_upvalue() {
        unsafe { upref.upvalue_value().value() }
    } else {
        upref
    };

    let name = if let Some(name) = unsafe { proto.upvalue_name((upvalue_index - 1) as usize) } {
        LuaString::from_interned(name)
    } else {
        LuaString::from_static(b"".as_bstr())
    };

    Some((name, value))
}

/// `append`
fn append_limited(output: &mut BString, data: &[u8], limit: usize) {
    if output.len() >= limit {
        return;
    }

    let remaining = limit - output.len();
    output.extend_from_slice(&data[..data.len().min(remaining)]);
}

/// `getcounters`
fn get_counters(
    thread: &Thread,
    proto: Proto,
    context: *mut (),
    function_visit: LuaCounterFunction,
    counter_visit: LuaCounterValue,
) {
    let get_counter_data = unsafe { thread.global() }.execution_counter_data();
    let proto_ref = unsafe { proto.as_ptr().as_ref().unwrap_unchecked() };
    if !proto_ref.exec_data.is_null()
        && let Some(get_counter_data) = get_counter_data
    {
        let mut count = 0usize;
        let mut data = unsafe { get_counter_data(thread, proto, &mut count) };

        if !data.is_null() && count != 0 {
            let debug_name_string = unsafe { proto.debug_name() };
            let debug_name = debug_name_string
                .as_ref()
                .map(|string| unsafe { string.as_bytes().as_bstr() });
            let line_defined = proto_ref.line_defined;

            function_visit(context, debug_name, line_defined);

            for _ in 0..count {
                let kind = unsafe { ptr::read_unaligned(data.cast::<u32>()) } as i32;
                data = unsafe { data.add(size_of::<u32>()) };

                let pc_pos = unsafe { ptr::read_unaligned(data.cast::<u32>()) };
                data = unsafe { data.add(size_of::<u32>()) };

                let hits = unsafe { ptr::read_unaligned(data.cast::<u64>()) };
                data = unsafe { data.add(size_of::<u64>()) };

                let line = if pc_pos == u32::MAX {
                    proto_ref.line_defined
                } else {
                    unsafe { proto.get_line(pc_pos as i32) }
                };

                counter_visit(context, kind, line, hits);
            }
        }
    }

    for index in 0..proto_ref.size_p as usize {
        let child = unsafe { proto.child_proto(index).unwrap_unchecked() };
        get_counters(thread, child, context, function_visit, counter_visit);
    }
}

/// `getmaxline`
fn get_max_line(proto: Proto) -> i32 {
    let mut result = -1;

    let proto_ref = unsafe { proto.as_ptr().as_ref().unwrap_unchecked() };
    for index in 0..proto_ref.size_code as usize {
        let line = unsafe { proto.get_line(index as i32) };
        result = result.max(line);
    }

    for index in 0..proto_ref.size_p as usize {
        let child = unsafe { proto.child_proto(index).unwrap_unchecked() };
        result = result.max(get_max_line(child));
    }

    result
}

/// `getnextline`
fn get_next_line(proto: Proto, line: i32) -> i32 {
    let mut closest = -1;

    let proto_ref = unsafe { proto.as_ptr().as_ref().unwrap_unchecked() };
    if !proto_ref.line_info.is_null() {
        for index in 0..proto_ref.size_code as usize {
            let instruction = Instruction::new(unsafe { *proto_ref.code.add(index) });
            if unsafe { instruction.opcode_unchecked() } == Opcode::PrepVarargs {
                continue;
            }

            let candidate = unsafe { proto.get_line(index as i32) };
            if candidate == line {
                return line;
            }

            if candidate > line && (closest == -1 || candidate < closest) {
                closest = candidate;
            }
        }
    }

    for index in 0..proto_ref.size_p as usize {
        let child = unsafe { proto.child_proto(index).unwrap_unchecked() };
        let candidate = get_next_line(child, line);
        if candidate == line {
            return line;
        }

        if candidate > line && (closest == -1 || candidate < closest) {
            closest = candidate;
        }
    }

    closest
}

/// `getcoverage`
fn get_coverage_recursive(
    proto: Proto,
    depth: i32,
    buffer: &mut [i32],
    context: *mut (),
    callback: LuaCoverage,
) {
    buffer.fill(-1);

    let proto_ref = unsafe { proto.as_ptr().as_ref().unwrap_unchecked() };
    for index in 0..proto_ref.size_code as usize {
        let instruction = Instruction::new(unsafe { *proto_ref.code.add(index) });
        if unsafe { instruction.opcode_unchecked() } != Opcode::Coverage {
            continue;
        }

        let line = unsafe { proto.get_line(index as i32) };
        let hits = instruction.e();
        debug_assert!((line as usize) < buffer.len());

        let entry = &mut buffer[line as usize];
        *entry = (*entry).max(hits);
    }

    let debug_name_string = unsafe { proto.debug_name() };
    let debug_name = debug_name_string
        .as_ref()
        .map(|string| unsafe { string.as_bytes().as_bstr() });
    let line_defined = proto_ref.line_defined;
    callback(context, debug_name, line_defined, depth, buffer);

    for index in 0..proto_ref.size_p as usize {
        let child = unsafe { proto.child_proto(index).unwrap_unchecked() };
        get_coverage_recursive(child, depth + 1, buffer, context, callback);
    }
}

impl Thread {
    /// `lua_callhook`
    pub unsafe fn call_hook<F>(&self, hook: F, userdata: *mut ()) -> crate::VmResult
    where
        F: FnOnce(&Thread, &mut LuaDebug) -> crate::VmResult,
    {
        debug_assert!(unsafe { self.current_call_info() != self.base_call_info() });
        unsafe { <Self as VmCallFrame>::call_hook(self, hook, userdata) }
    }

    /// `lua_stackdepth`
    pub unsafe fn stack_depth(&self) -> i32 {
        unsafe {
            self.current_call_info_cursor()
                .offset_from(self.base_call_info_cursor()) as i32
        }
    }

    /// `lua_getinfo`
    pub unsafe fn get_info(&self, level: i32, what: &str, ar: &mut LuaDebug) -> VmErrorResult<i32> {
        unsafe {
            let mut closure = None;
            let mut call_info = None;
            let mut stack_function = None;

            if level < 0 {
                let available = self.stack_top().offset_from(self.stack_base()) as i32;
                if -level > available {
                    return Ok(0);
                }

                let function = self.stack_top().offset(level as isize);
                if !function.value_unchecked().is_function() {
                    return Ok(0);
                }

                stack_function = Some(function);
                closure = Some(function.value_unchecked().closure_value());
            } else if (level as usize)
                < self
                    .current_call_info_cursor()
                    .offset_from(self.base_call_info_cursor()) as usize
            {
                let ci = self
                    .current_call_info_cursor()
                    .sub(level as usize)
                    .call_info_unchecked();
                call_info = Some(ci);
                closure = Some(ci.function_closure());
            }

            let Some(closure) = closure else {
                return Ok(0);
            };

            if aux_get_info(
                what.as_bytes(),
                ar,
                closure,
                call_info.map(|ci| ci.as_ptr()),
            ) {
                self.thread_barrier();

                if let Some(call_info) = call_info {
                    self.push_value_internal(call_info.function().value_unchecked())?;
                } else {
                    self.push_value_internal(stack_function.unwrap_unchecked().value_unchecked())?;
                }
            }

            Ok(1)
        }
    }

    /// `lua_getargument`
    pub unsafe fn get_argument(&self, level: i32, argument: i32) -> VmErrorResult<i32> {
        unsafe {
            let depth = self.stack_depth();
            if (level as u32) >= (depth as u32) {
                return Ok(0);
            }

            let call_info = self
                .current_call_info_cursor()
                .sub(level as usize)
                .call_info_unchecked();
            if call_info.as_ptr().as_ref().unwrap_unchecked().flags
                & crate::state::LUA_CALLINFO_NATIVE
                != 0
            {
                return Ok(0);
            }

            let Some(proto) = lua_proto(call_info) else {
                return Ok(0);
            };

            if argument <= 0 {
                return Ok(0);
            }

            if argument <= i32::from(proto.as_ptr().as_ref().unwrap_unchecked().num_params) {
                self.thread_barrier();
                self.push_value_internal(
                    call_info
                        .base()
                        .add((argument - 1) as usize)
                        .value_unchecked(),
                )?;
                Ok(1)
            } else if proto.as_ptr().as_ref().unwrap_unchecked().is_vararg != 0
                && argument < call_info.base().offset_from(call_info.function()) as i32
            {
                self.thread_barrier();
                self.push_value_internal(
                    call_info
                        .function()
                        .add(argument as usize)
                        .value_unchecked(),
                )?;
                Ok(1)
            } else {
                Ok(0)
            }
        }
    }

    /// `lua_getlocal`
    pub unsafe fn get_local(&self, level: i32, local: i32) -> VmErrorResult<Option<LuaString>> {
        unsafe {
            let depth = self.stack_depth();
            if (level as u32) >= (depth as u32) {
                return Ok(None);
            }

            let call_info = self
                .current_call_info_cursor()
                .sub(level as usize)
                .call_info_unchecked();
            if call_info.as_ptr().as_ref().unwrap_unchecked().flags
                & crate::state::LUA_CALLINFO_NATIVE
                != 0
            {
                return Ok(None);
            }

            let Some(proto) = lua_proto(call_info) else {
                return Ok(None);
            };
            let Some(var) = proto.get_local(local, current_pc(call_info, proto)) else {
                return Ok(None);
            };

            self.thread_barrier();
            self.push_value_internal(
                call_info
                    .base()
                    .add(var.as_ptr().as_ref().unwrap_unchecked().reg as usize)
                    .value_unchecked(),
            )?;

            let name = var.name();
            debug_assert!(name.is_some());
            Ok(Some(LuaString::from_interned(name.unwrap_unchecked())))
        }
    }

    /// `lua_setlocal`
    pub unsafe fn set_local(&self, level: i32, local: i32) -> Option<LuaString> {
        unsafe {
            debug_assert!(self.stack_top() > self.stack_base());

            let depth = self.stack_depth();
            if (level as u32) >= (depth as u32) {
                return None;
            }

            let call_info = self
                .current_call_info_cursor()
                .sub(level as usize)
                .call_info_unchecked();
            if call_info.as_ptr().as_ref().unwrap_unchecked().flags
                & crate::state::LUA_CALLINFO_NATIVE
                != 0
            {
                return None;
            }

            let proto = lua_proto(call_info)?;
            let var = proto.get_local(local, current_pc(call_info, proto))?;
            call_info
                .base()
                .add(var.as_ptr().as_ref().unwrap_unchecked().reg as usize)
                .value_unchecked()
                .set_obj(self.stack_top().sub(1).value_unchecked());
            self.set_stack_top(self.stack_top().sub(1));

            let name = var.name();
            debug_assert!(name.is_some());
            Some(LuaString::from_interned(name.unwrap_unchecked()))
        }
    }

    /// `lua_getupvalue`
    pub unsafe fn get_upvalue(
        &self,
        function_index: i32,
        upvalue_index: i32,
    ) -> VmErrorResult<Option<LuaString>> {
        unsafe {
            self.thread_barrier();
            self.ensure_stack(self, 1)?;

            let function = self.to_object(function_index).unwrap_unchecked();
            let Some((name, value)) = aux_upvalue(function, upvalue_index) else {
                return Ok(None);
            };

            self.push_value_internal(value)?;
            Ok(Some(name))
        }
    }

    /// `lua_setupvalue`
    pub unsafe fn set_upvalue(&self, function_index: i32, upvalue_index: i32) -> Option<LuaString> {
        unsafe {
            let function = self.to_object(function_index).unwrap_unchecked();
            let (name, value) = aux_upvalue(function, upvalue_index)?;

            self.set_stack_top(self.stack_top().sub(1));
            value.set_obj(self.stack_top().value_unchecked());

            let closure = function.closure_value();
            let written = self.stack_top().value_unchecked();
            if written.is_collectable() {
                let object = closure.into();
                let child = written.gc_value();
                self.barrier_forward(object, child);
            }

            Some(name)
        }
    }

    /// `lua_singlestep`
    pub unsafe fn single_step(&self, enabled: i32) {
        unsafe { self.as_ptr().as_mut().unwrap_unchecked().single_step = enabled != 0 };
    }

    /// `lua_breakpoint`
    pub unsafe fn breakpoint(
        &self,
        function_index: i32,
        line: i32,
        enabled: i32,
    ) -> VmErrorResult<i32> {
        unsafe {
            let function = self.to_object(function_index).unwrap_unchecked();
            debug_assert!(function.is_function());

            let closure = function.closure_value();
            debug_assert!(closure.is_lua());

            let proto = closure.proto().unwrap_unchecked();

            let target = get_next_line(proto, line);
            if target != -1 {
                self.breakpoint_internal(proto, target, enabled != 0)?;
            }

            Ok(target)
        }
    }

    /// `lua_hascustomexecution`
    pub unsafe fn has_custom_execution(&self, level: i32) -> i32 {
        unsafe { self.has_native(level) }
    }

    /// `lua_incustomexecution`
    pub unsafe fn in_custom_execution(&self, level: i32) -> i32 {
        unsafe { self.is_native(level) }
    }

    /// `lua_atbreakpoint`
    pub unsafe fn at_breakpoint(&self) -> i32 {
        i32::from(unsafe { self.on_break() })
    }

    /// `lua_getcoverage`
    pub unsafe fn get_coverage(
        &self,
        function_index: i32,
        context: *mut (),
        callback: LuaCoverage,
    ) -> VmErrorResult {
        unsafe {
            let function = self.to_object(function_index).unwrap_unchecked();
            debug_assert!(function.is_function());

            let closure = function.closure_value();
            debug_assert!(closure.is_lua());

            let proto = closure.proto().unwrap_unchecked();

            let size = get_max_line(proto) + 1;
            if size <= 0 {
                return Ok(());
            }

            let buffer = self.new_array::<i32>(size as usize, 0)?;
            let buffer = core::slice::from_raw_parts_mut(buffer, size as usize);

            get_coverage_recursive(proto, 0, buffer, context, callback);

            self.free_array(buffer.as_mut_ptr(), size as usize, 0);
        }
        Ok(())
    }

    /// `lua_getcounters`
    pub unsafe fn get_counters(
        &self,
        function_index: i32,
        context: *mut (),
        function_visit: LuaCounterFunction,
        counter_visit: LuaCounterValue,
    ) {
        unsafe {
            let Some(function) = self.to_object(function_index) else {
                return;
            };
            if !function.is_function() {
                return;
            }

            if self.global().execution_counter_data().is_none() {
                return;
            }

            let closure = function.closure_value();
            if closure.is_native() {
                return;
            }

            let proto = closure.proto().unwrap_unchecked();
            get_counters(self, proto, context, function_visit, counter_visit);
        }
    }

    /// `lua_debugtrace`
    pub unsafe fn debug_trace(&self) -> VmErrorResult<BString> {
        const BUFFER_SIZE: usize = 4096;
        const LIMIT_1: i32 = 10;
        const LIMIT_2: i32 = 10;
        const BYTE_LIMIT: usize = BUFFER_SIZE - 1;

        unsafe {
            let depth = self.stack_depth();
            let mut output = BString::new(Vec::new());
            let mut ar = LuaDebug::default();
            let mut level = 0;

            while self.get_info(level, "sln", &mut ar)? != 0 {
                if !ar.source.as_bytes().is_empty() {
                    append_limited(&mut output, ar.short_src(), BYTE_LIMIT);
                }

                if ar.currentline > 0 {
                    append_limited(&mut output, b":", BYTE_LIMIT);
                    let mut line = [0u8; crate::number::LUAI_MAXINT2STR];
                    let len = crate::number::int_to_str(&mut line, i64::from(ar.currentline));
                    append_limited(&mut output, &line[..len], BYTE_LIMIT);
                }

                if let Some(name) = ar.name {
                    append_limited(&mut output, b" function ", BYTE_LIMIT);
                    append_limited(&mut output, name.as_bytes(), BYTE_LIMIT);
                }

                append_limited(&mut output, b"\n", BYTE_LIMIT);

                if depth > LIMIT_1 + LIMIT_2 && level == LIMIT_1 - 1 {
                    append_limited(&mut output, b"... (+", BYTE_LIMIT);
                    let mut skipped = [0u8; crate::number::LUAI_MAXINT2STR];
                    let skipped_frames = i64::from(depth - LIMIT_1 - LIMIT_2);
                    let len = crate::number::int_to_str(&mut skipped, skipped_frames);
                    append_limited(&mut output, &skipped[..len], BYTE_LIMIT);
                    append_limited(&mut output, b" frames)\n", BYTE_LIMIT);
                    level = depth - LIMIT_2 - 1;
                }

                level += 1;
                ar = LuaDebug::default();
            }

            Ok(output)
        }
    }
}

// Traceback construction
impl Thread {
    /// `luaL_traceback`
    pub unsafe fn traceback(
        &self,
        source: Option<&Thread>,
        message: Option<&BStr>,
        level: i32,
    ) -> VmErrorResult {
        debug_assert!(level >= 0);

        let source = source.unwrap_or(self);
        unsafe {
            let mut buffer_storage = LuaStringBuilderStorage::uninit();
            let mut buffer = LuaStringBuilder::new(self, &mut buffer_storage);

            if let Some(message) = message {
                buffer.push_bytes(message.as_bytes())?;
                buffer.push_byte(b'\n')?;
            }

            let mut frame = level;
            let mut ar = LuaDebug::default();
            while source.get_info(frame, "sln", &mut ar)? != 0 {
                frame += 1;

                if ar.what.as_bytes() == b"C" {
                    ar = LuaDebug::default();
                    continue;
                }

                if !ar.source.as_bytes().is_empty() {
                    buffer.push_bytes(ar.short_src())?;
                }

                let mut line_digits = None;
                if ar.currentline > 0 {
                    let mut digits = [0u8; 32];
                    let mut end = digits.len();
                    let mut line = ar.currentline as u32;
                    while line > 0 {
                        end -= 1;
                        digits[end] = b'0' + (line % 10) as u8;
                        line /= 10;
                    }
                    line_digits = Some((digits, end));
                }

                if let Some((digits, end)) = line_digits {
                    buffer.push_byte(b':')?;
                    buffer.push_bytes(&digits[end..])?;
                }

                if let Some(name) = ar.name {
                    buffer.push_bytes(b" function ")?;
                    buffer.push_bytes(name.as_bytes())?;
                }

                buffer.push_byte(b'\n')?;
                ar = LuaDebug::default();
            }

            buffer.finish()?;
        }
        Ok(())
    }
}