luau 0.732.0

use luau_vm::thread::{
    LUA_COERR, LUA_COFIN, LUA_CONOR, LUA_CORUN, LUA_COSUS, LUA_GLOBALS_INDEX, LUA_MULTRET,
    StackGuard,
};
use luau_vm::{Thread as VmThread, VmControl, VmError, VmExit};

use super::Thread;
use crate::error::Error;
use crate::function::Function;
use crate::lua::{ChunkLoad, Lua, LuaRef};
use crate::table::Table;
use crate::value::{FromLuaMulti, IntoLua, IntoLuaMulti};

/// The execution status of a Luau thread.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ThreadStatus {
    /// The thread is new or suspended and can be resumed.
    Resumable,
    /// The thread is currently running.
    Running,
    /// The thread resumed another thread that is currently running.
    Normal,
    /// The thread completed successfully.
    Finished,
    /// The thread stopped with an error.
    Error,
}

impl Lua {
    /// Creates a coroutine from a function.
    pub fn create_thread<'lua>(
        &'lua self,
        function: Function<'lua>,
    ) -> Result<Thread<'lua>, Error> {
        self.lua_ref().create_thread(function)
    }
}

impl<'lua> LuaRef<'lua> {
    /// Creates a coroutine from a function.
    pub fn create_thread(&self, function: Function<'lua>) -> Result<Thread<'lua>, Error> {
        if !unsafe { self.as_vm().same_vm(function.thread().as_vm()) } {
            return Err(Error::foreign_lua_handle());
        }
        function.into_thread()
    }
}

impl<'lua> Thread<'lua> {
    pub(crate) fn create_empty_thread(&self) -> Result<Thread<'lua>, Error> {
        unsafe {
            let vm_thread = self.as_vm();
            let _stack = StackGuard::new(vm_thread);
            let thread = vm_thread
                .new_thread()
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            let pointer = vm_thread.to_pointer(-1);
            let reference = vm_thread
                .ref_value(-1)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;
            Ok(Thread::from_ref(
                self.reference_thread(),
                self.runtime(),
                thread,
                reference,
                pointer,
            ))
        }
    }

    /// Returns this thread's execution status.
    pub fn status(&self) -> ThreadStatus {
        match unsafe { self.reference_thread().co_status(self.as_vm()) } {
            LUA_CORUN => ThreadStatus::Running,
            LUA_COSUS => ThreadStatus::Resumable,
            LUA_CONOR => ThreadStatus::Normal,
            LUA_COFIN => ThreadStatus::Finished,
            LUA_COERR => ThreadStatus::Error,
            status => unreachable!("unexpected coroutine status {status}"),
        }
    }

    /// Returns whether this thread can be resumed.
    pub fn is_resumable(&self) -> bool {
        self.status() == ThreadStatus::Resumable
    }

    /// Returns whether this thread is currently running.
    pub fn is_running(&self) -> bool {
        self.status() == ThreadStatus::Running
    }

    /// Returns whether this thread resumed another active thread.
    pub fn is_normal(&self) -> bool {
        self.status() == ThreadStatus::Normal
    }

    /// Returns whether this thread completed successfully.
    pub fn is_finished(&self) -> bool {
        self.status() == ThreadStatus::Finished
    }

    /// Returns whether this thread stopped with an error.
    pub fn is_error(&self) -> bool {
        self.status() == ThreadStatus::Error
    }

    /// Resets this coroutine and replaces its body.
    ///
    /// Running and normal threads cannot be reset.
    pub fn reset(&self, function: Function<'lua>) -> Result<(), Error> {
        if matches!(self.status(), ThreadStatus::Running | ThreadStatus::Normal) {
            return Err(Error::runtime("cannot reset an active coroutine"));
        }
        if !self.same_vm(function.thread().as_vm()) {
            return Err(Error::foreign_lua_handle());
        }

        unsafe {
            self.as_vm()
                .reset()
                .map_err(|error| Error::from_thread_exit(self.as_vm(), error))?;
        }
        function.push_to(self)
    }

    pub(crate) fn load_body(
        &self,
        chunk_name: impl AsRef<[u8]>,
        bytecode: impl AsRef<[u8]>,
        environment: Option<&Table<'_>>,
        load: ChunkLoad,
    ) -> Result<(), Error> {
        unsafe {
            let thread = self.as_vm();
            if let Err(error) =
                self.load_chunk(chunk_name.as_ref(), bytecode.as_ref(), environment, load)
            {
                thread.restore_top(0);
                return Err(error);
            }
        }
        Ok(())
    }

    /// Resumes this coroutine with the given arguments.
    ///
    /// Returns the values yielded or returned by the coroutine.
    pub fn resume<'thread, R>(&'thread self, args: impl IntoLuaMulti<'thread>) -> Result<R, Error>
    where
        R: FromLuaMulti<'thread>,
    {
        if !self.is_resumable() {
            return Err(Error::CoroutineUnresumable);
        }

        unsafe {
            let vm_thread = self.as_vm();
            let stack = StackGuard::new(vm_thread);
            let arg_count =
                i32::try_from(args.push_into_stack_multi(self)?).map_err(|_| Error::StackError)?;

            let resumed = self
                .runtime()
                .with_current_thread(|from| vm_thread.resume(Some(from), arg_count));
            match resumed {
                Ok(()) | Err(VmExit::Control(VmControl::Yield)) => {}
                Err(VmExit::Error(error)) => {
                    let error = error_with_traceback(vm_thread, error);
                    stack.dismiss();
                    vm_thread.restore_top(0);
                    return Err(error);
                }
                Err(VmExit::Control(VmControl::Break)) => {
                    stack.dismiss();
                    vm_thread.restore_top(0);
                    return Err(Error::Interrupted);
                }
            }

            stack.dismiss();
            let result_count = vm_thread.get_top();
            let result = R::from_stack_multi(self, 0, result_count);
            vm_thread.restore_top(0);
            result
        }
    }

    /// Resumes this coroutine by raising `error` at its suspended point.
    pub fn resume_error<'thread, R>(&'thread self, error: impl IntoLua<'thread>) -> Result<R, Error>
    where
        R: FromLuaMulti<'thread>,
    {
        if !self.is_resumable() {
            return Err(Error::CoroutineUnresumable);
        }

        unsafe {
            let vm_thread = self.as_vm();
            let stack = StackGuard::new(vm_thread);
            error.push_into_stack(self)?;

            let resumed = self
                .runtime()
                .with_current_thread(|from| vm_thread.resume_error(Some(from)));
            match resumed {
                Ok(()) | Err(VmExit::Control(VmControl::Yield)) => {}
                Err(VmExit::Error(error)) => {
                    let error = error_with_traceback(vm_thread, error);
                    stack.dismiss();
                    vm_thread.restore_top(0);
                    return Err(error);
                }
                Err(VmExit::Control(VmControl::Break)) => {
                    stack.dismiss();
                    vm_thread.restore_top(0);
                    return Err(Error::Interrupted);
                }
            }

            stack.dismiss();
            let result_count = vm_thread.get_top();
            let result = R::from_stack_multi(self, 0, result_count);
            vm_thread.restore_top(0);
            result
        }
    }

    pub(crate) fn load_bytecode(
        &self,
        chunk_name: impl AsRef<[u8]>,
        bytecode: impl AsRef<[u8]>,
        environment: Option<&Table<'_>>,
        load: ChunkLoad,
    ) -> Result<Function<'lua>, Error> {
        unsafe {
            let thread = self.as_vm();
            let _stack = StackGuard::new(thread);
            self.load_chunk(chunk_name.as_ref(), bytecode.as_ref(), environment, load)?;

            let pointer = thread.to_pointer(-1);
            let reference = match thread.ref_value(-1) {
                Ok(reference) => reference,
                Err(exit) => return Err(Error::from_thread_exit(thread, exit)),
            };

            Ok(Function::from_ref(
                self.reference_thread(),
                self.runtime(),
                reference,
                pointer,
            ))
        }
    }

    pub(crate) fn run_bytecode_with_args<R>(
        &self,
        chunk_name: impl AsRef<[u8]>,
        bytecode: impl AsRef<[u8]>,
        args: impl IntoLuaMulti<'lua>,
        environment: Option<&Table<'_>>,
        load: ChunkLoad,
    ) -> Result<R, Error>
    where
        R: FromLuaMulti<'lua>,
    {
        unsafe {
            let vm_thread = self.as_vm();
            let stack = StackGuard::new(vm_thread);
            self.load_chunk(chunk_name.as_ref(), bytecode.as_ref(), environment, load)?;
            let arg_count =
                i32::try_from(args.push_into_stack_multi(self)?).map_err(|_| Error::StackError)?;

            vm_thread
                .protected_call(arg_count, LUA_MULTRET, 0)
                .map_err(|exit| Error::from_thread_exit(vm_thread, exit))?;

            let result_count = vm_thread.get_top() - stack.top();
            R::from_stack_multi(self, stack.top(), result_count)
        }
    }

    fn load_chunk(
        &self,
        chunk_name: &[u8],
        bytecode: &[u8],
        environment: Option<&Table<'_>>,
        load: ChunkLoad,
    ) -> Result<(), Error> {
        unsafe {
            let thread = self.as_vm();
            let environment_index = match environment {
                Some(environment) => {
                    environment.push_to(self)?;
                    thread.get_top()
                }
                None => {
                    thread
                        .push_value(LUA_GLOBALS_INDEX)
                        .map_err(|error| Error::from_thread_exit(thread, error))?;
                    thread.get_top()
                }
            };
            let uses_current_environment =
                thread.raw_equal(environment_index, LUA_GLOBALS_INDEX) != 0;

            load.before_load(
                self.runtime(),
                thread,
                environment_index,
                uses_current_environment,
            );
            thread
                .load(chunk_name, bytecode, environment_index)
                .map_err(|exit| Error::from_thread_exit(thread, exit))?;
            load.loaded(self.runtime(), uses_current_environment);
            thread.remove(environment_index);
        }

        Ok(())
    }
}

fn error_with_traceback(thread: &VmThread, error: VmError) -> Error {
    let mut error = Error::from_thread_exit(thread, error);
    let Error::RuntimeError(message) = &mut error else {
        return error;
    };
    let Ok(traceback) = (unsafe { thread.debug_trace() }) else {
        return error;
    };
    let traceback = String::from_utf8_lossy(traceback.as_slice());
    let traceback = traceback.trim();
    if !traceback.is_empty() {
        message.push_str("\nstack traceback:\n");
        message.push_str(traceback);
    }
    error
}