luau-vm 0.732.0

use luau_bytecode::opcodes::PROTO_FLAG_USES_EXPORT;
use luau_common::{ByteSlice, flags};

use super::stack::RawStackAccess;
use super::{LUA_MULTRET, Thread};
use crate::call::{CallRuntime, ProtectedCall, ThreadStack, resume_handle, resume_protected};
use crate::debug::DebugRuntime;
use crate::function::{FunctionRuntime, RawNativeClosure};
use crate::gc::{GcBarrier, GcRuntime};
use crate::handle::RawHandle;
use crate::native::{
    NativeCallContext, NativeCallResult, RawNativeContinuation, RawNativeFunction,
};
use crate::state::{
    LUA_CALLINFO_HANDLE, LUA_OK, THREAD_STATUS_BREAK, THREAD_STATUS_OK, THREAD_STATUS_YIELD,
    ThreadState,
};
use crate::string::StringRuntime;
use crate::value::{TValueCursor, nil_object};
use crate::{VmControl, VmError, VmErrorResult, VmExit, VmResult};

// Callable values
impl Thread {
    /// `lua_pushcclosurek`
    pub unsafe fn push_native_closure_k(
        &self,
        function: RawNativeFunction,
        debug_name: Option<&'static str>,
        upvalues: i32,
        continuation: Option<RawNativeContinuation>,
    ) -> VmErrorResult {
        debug_assert!(upvalues >= 0);
        debug_assert!(
            upvalues <= unsafe { self.stack_top().offset_from(self.stack_base()) as i32 }
        );
        unsafe {
            self.check_gc()?;
            self.thread_barrier();
            self.ensure_stack(self, 1)?;

            let closure = self.new_native_closure(upvalues, Some(self.current_env()))?;
            let managed_debug_name = if flags::LuauManagedDebugNames.get() {
                debug_name
                    .map(|name| self.intern_string(name.as_bytes().as_bstr()))
                    .transpose()?
            } else {
                None
            };
            closure.set_native_data(RawNativeClosure {
                function: Some(function),
                continuation,
                debug_name_deprecated: (!flags::LuauManagedDebugNames.get())
                    .then_some(debug_name)
                    .flatten(),
                debug_name: managed_debug_name.map_or(core::ptr::null_mut(), |name| name.as_ptr()),
            });
            self.set_stack_top(self.stack_top().sub(upvalues as usize));

            for index in (0..upvalues as usize).rev() {
                closure
                    .native_upvalue(index)
                    .set_obj(self.stack_top().add(index).value_unchecked());
            }

            let top = self.stack_top();
            top.value_unchecked().set_closure_value(closure);
            debug_assert!(top < self.current_call_info().top());
            self.set_stack_top(top.add(1));
        }
        Ok(())
    }

    /// `lua_pushcfunction`
    pub unsafe fn push_native_function(
        &self,
        function: RawNativeFunction,
        debug_name: Option<&'static str>,
    ) -> VmErrorResult {
        unsafe { self.push_native_closure_k(function, debug_name, 0, None) }
    }

    /// `lua_pushcclosure`
    pub unsafe fn push_native_closure(
        &self,
        function: RawNativeFunction,
        debug_name: Option<&'static str>,
        upvalues: i32,
    ) -> VmErrorResult {
        unsafe { self.push_native_closure_k(function, debug_name, upvalues, None) }
    }

    /// `lua_tocfunction`
    pub unsafe fn to_native_function(&self, index: i32) -> Option<RawNativeFunction> {
        unsafe {
            let object = self.index_to_addr(index);
            if object == nil_object() || !object.is_function() {
                return None;
            }

            let closure = object.closure_value();
            if !closure.is_native() {
                return None;
            }

            closure.native_data().function
        }
    }

    /// `lua_clonefunction`
    pub unsafe fn clone_function(&self, index: i32) -> VmErrorResult {
        unsafe {
            self.check_gc()?;
            self.thread_barrier();
            self.ensure_stack(self, 1)?;

            let object = self.to_object(index).unwrap_unchecked();
            debug_assert!(object.is_function());

            let closure = object.closure_value();
            let closure_ref = closure.as_ptr().as_ref().unwrap_unchecked();
            debug_assert!(closure.is_lua());

            let proto = closure.proto().unwrap_unchecked();
            let environment = self.globals();
            let new_closure =
                self.new_lua_closure(i32::from(closure_ref.n_upvalues), Some(environment), proto)?;

            for index in 0..closure_ref.n_upvalues as usize {
                new_closure
                    .lua_upvalue_ref(index)
                    .set_obj(closure.lua_upvalue_ref(index));
            }

            let top = self.stack_top();
            top.value_unchecked().set_closure_value(new_closure);
            debug_assert!(top < self.current_call_info().top());
            self.set_stack_top(top.add(1));
        }
        Ok(())
    }

    /// `lua_usesexport`
    pub unsafe fn uses_export(&self, index: i32) -> i32 {
        let object = unsafe { self.index_to_addr(index) };
        if object == nil_object() || !object.is_function() {
            return 0;
        }

        let closure = object.closure_value();
        if unsafe { closure.is_native() } {
            return 0;
        }

        let proto = unsafe { closure.proto().unwrap_unchecked() };
        i32::from(
            unsafe { proto.as_ptr().as_ref().unwrap_unchecked().flags } & PROTO_FLAG_USES_EXPORT
                != 0,
        )
    }
}

// Calls and protected calls
/// `adjustresults`
fn adjust_results(thread: &Thread, n_results: i32) {
    if n_results == LUA_MULTRET && unsafe { thread.stack_top() >= thread.current_call_info().top() }
    {
        unsafe {
            thread.current_call_info().set_top(thread.stack_top());
        }
    }
}

/// `checkresults`
fn check_results(thread: &Thread, n_args: i32, n_results: i32) {
    debug_assert!(
        n_results == LUA_MULTRET
            || unsafe {
                thread
                    .current_call_info()
                    .top()
                    .offset_from(thread.stack_top()) as i32
            } >= n_results - n_args
    );
}

impl Thread {
    /// `lua_call`
    pub unsafe fn call(&self, n_args: i32, n_results: i32) -> VmResult {
        debug_assert!(n_args >= 0);
        debug_assert!(n_results >= LUA_MULTRET);
        debug_assert!(n_args < unsafe { self.stack_top().offset_from(self.stack_base()) as i32 });
        if n_results > n_args + 1 {
            unsafe { self.ensure_stack(self, n_results - (n_args + 1))? };
        }
        check_results(self, n_args, n_results);

        let result = unsafe {
            debug_assert!(self.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_OK);
            let function = self.stack_top().sub((n_args + 1) as usize);
            self.call_internal(function, n_results)
        };

        if matches!(result, Ok(()) | Err(VmExit::Control(_))) {
            adjust_results(self, n_results);
        }

        result
    }

    /// `lua_pcall`
    pub unsafe fn protected_call(&self, n_args: i32, n_results: i32, errfunc: i32) -> VmResult {
        debug_assert!(n_args >= 0);
        debug_assert!(n_results >= LUA_MULTRET);
        #[repr(C)]
        struct CallContext {
            function: TValueCursor,
            n_results: i32,
        }

        /// `f_call`
        unsafe fn pcall_callback(thread: &Thread, context: &mut CallContext) -> VmResult {
            unsafe { thread.call_internal(context.function, context.n_results)? };
            Ok(())
        }

        if n_results > n_args + 1
            && let Err(error) = unsafe { self.ensure_stack(self, n_results - (n_args + 1)) }
        {
            return Err(error.into());
        }
        check_results(self, n_args, n_results);

        unsafe {
            debug_assert!(n_args < self.stack_top().offset_from(self.stack_base()) as i32);
            debug_assert!(self.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_OK);

            let error_function = if errfunc != 0 {
                let error_function = self.stack_index_to_cursor(errfunc);
                debug_assert!(error_function.value_unchecked() != nil_object());
                self.save_stack(error_function)
            } else {
                0
            };

            let function = self.stack_top().sub((n_args + 1) as usize);
            let mut context = CallContext {
                function,
                n_results,
            };

            let result = self.protected_call_internal(
                pcall_callback,
                &mut context,
                self.save_stack(function),
                error_function,
            );
            adjust_results(self, n_results);
            result
        }
    }

    /// `lua_cpcall`
    pub unsafe fn protected_native_call(
        &self,
        function: RawNativeFunction,
        userdata: *mut (),
    ) -> VmResult {
        #[repr(C)]
        struct ProtectedCallContext {
            function: RawNativeFunction,
            userdata: *mut (),
        }

        /// `f_Ccall`
        unsafe fn protected_call_callback(
            thread: &Thread,
            context: &mut ProtectedCallContext,
        ) -> VmResult {
            unsafe {
                if thread.check_stack(2) == 0 {
                    return crate::run_error!(thread, "stack limit").map_err(Into::into);
                }

                thread.push_native_closure_k(context.function, None, 0, None)?;
                thread.push_light_userdata_tagged(context.userdata, 0)?;
                let function = thread.stack_top().sub(2);
                thread.call_internal(function, 0)?;
            }
            Ok(())
        }

        let mut context = ProtectedCallContext { function, userdata };
        unsafe {
            debug_assert!(self.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_OK);
            self.protected_call_internal(
                protected_call_callback,
                &mut context,
                self.save_stack(self.stack_top()),
                0,
            )
        }
    }

    /// `lua_error`
    pub unsafe fn error<T>(&self) -> VmErrorResult<T> {
        debug_assert!(unsafe { self.stack_top() > self.stack_base() });
        Err(VmError::Runtime)
    }

    /// `lua_yield`
    #[must_use = "yield_current returns VM control flow that must be propagated by the caller"]
    pub unsafe fn yield_current(&self, results: i32) -> NativeCallResult {
        unsafe {
            debug_assert!(results >= 0);
            debug_assert!(results <= self.stack_top().offset_from(self.stack_base()) as i32);

            if self.as_ptr().as_ref().unwrap_unchecked().native_call_depth
                > self
                    .as_ptr()
                    .as_ref()
                    .unwrap_unchecked()
                    .base_native_call_depth
            {
                return crate::run_error!(
                    self,
                    "attempt to yield across metamethod/C-call boundary"
                )
                .map_err(Into::into);
            }

            self.set_stack_base(self.stack_top().sub(results as usize));
            self.as_ptr().as_mut().unwrap_unchecked().status = THREAD_STATUS_YIELD;
        }

        Err(VmExit::Control(VmControl::Yield))
    }

    /// `lua_break`
    #[must_use = "break_current returns VM control flow that must be propagated by the caller"]
    pub unsafe fn break_current(&self) -> NativeCallResult {
        unsafe {
            if self.as_ptr().as_ref().unwrap_unchecked().native_call_depth
                > self
                    .as_ptr()
                    .as_ref()
                    .unwrap_unchecked()
                    .base_native_call_depth
            {
                return crate::run_error!(
                    self,
                    "attempt to break across metamethod/C-call boundary"
                )
                .map_err(Into::into);
            }

            self.as_ptr().as_mut().unwrap_unchecked().status = THREAD_STATUS_BREAK;
        }

        Err(VmExit::Control(VmControl::Break))
    }
}

// Coroutine resumption
impl Thread {
    /// `lua_resume`
    pub unsafe fn resume(&self, from: Option<&Thread>, nargs: i32) -> VmResult {
        unsafe {
            self.resume_start(from, nargs)?;

            let old_native_call_depth = self.as_ptr().as_ref().unwrap_unchecked().native_call_depth;
            let mut first_argument = self.stack_top().sub(nargs as usize);
            let result = self.raw_run_protected(resume_protected, &mut first_argument);

            self.resume_finish(result, old_native_call_depth)
        }
    }

    /// `lua_resumeerror`
    pub unsafe fn resume_error(&self, from: Option<&Thread>) -> VmResult {
        unsafe {
            self.resume_start(from, 1)?;

            let old_native_call_depth = self.as_ptr().as_ref().unwrap_unchecked().native_call_depth;
            let mut result = Err(VmError::Runtime.into());

            if let Some(handler) = self.resume_find_handler() {
                self.as_ptr().as_mut().unwrap_unchecked().status = VmError::Runtime.status() as u8;
                let mut handler = handler;
                result = self.raw_run_protected(resume_handle, &mut handler);
            }

            self.resume_finish(result, old_native_call_depth)
        }
    }
}

// Yieldable native calls
impl Thread {
    /// `luaL_callyieldable`
    pub unsafe fn call_yieldable(&self, nargs: i32, nresults: i32) -> NativeCallResult {
        unsafe {
            let closure = self.current_function();
            debug_assert!(closure.is_native());
            debug_assert!(closure.native_data().continuation.is_some());

            self.call(nargs, nresults)?;
            closure.native_data().continuation.unwrap_unchecked()(
                NativeCallContext::new(self),
                LUA_OK,
            )
        }
    }

    /// `luaL_pcallyieldable`
    pub unsafe fn protected_call_yieldable(
        &self,
        nargs: i32,
        nresults: i32,
        errfunc: i32,
    ) -> NativeCallResult {
        #[repr(C)]
        struct CallContext {
            function: TValueCursor,
            n_results: i32,
        }

        unsafe fn protected_call_yieldable_run(
            thread: &Thread,
            context: &mut CallContext,
        ) -> VmResult {
            unsafe {
                thread.call_int(
                    context.function,
                    context.n_results,
                    thread.is_yieldable() != 0,
                )?;
            }
            Ok(())
        }

        unsafe {
            let closure = self.current_function();
            debug_assert!(closure.is_native());
            let continuation = closure.native_data().continuation.unwrap_unchecked();
            debug_assert!(nargs < self.stack_top().offset_from(self.stack_base()) as i32);
            debug_assert!(errfunc >= 0);
            debug_assert!(errfunc <= self.stack_top().offset_from(self.stack_base()) as i32);

            let call_info = self.current_call_info();
            call_info.set_errfunc(errfunc);
            call_info.as_ptr().as_mut().unwrap_unchecked().flags |= LUA_CALLINFO_HANDLE;

            let function = self.stack_top().sub((nargs + 1) as usize);
            let mut context = CallContext {
                function,
                n_results: nresults,
            };
            let saved_function = self.save_stack(function);
            let saved_error_function = if errfunc != 0 {
                self.save_stack(self.stack_base().add((errfunc - 1) as usize))
            } else {
                0
            };

            let result = self.protected_call_internal(
                protected_call_yieldable_run,
                &mut context,
                saved_function,
                saved_error_function,
            );

            self.expand_stack_limit(self.stack_top());

            let status = match result {
                Ok(()) => LUA_OK,
                Err(VmExit::Error(error)) => error.status(),
                Err(exit) => return Err(exit),
            };

            call_info.as_ptr().as_mut().unwrap_unchecked().flags &= !LUA_CALLINFO_HANDLE;

            continuation(NativeCallContext::new(self), status)
        }
    }
}