luau-vm 0.732.0

Pure-Rust Luau virtual machine, garbage collector, and standard libraries
Documentation
use crate::handle::RawHandle;
use crate::{VmControl, VmExit, VmResult};

use crate::native::{NativeCallContext, NativeCallResult, NativeFunction};
use crate::state::{
    THREAD_STATUS_BREAK, THREAD_STATUS_ERR_ERR, THREAD_STATUS_ERR_MEM, THREAD_STATUS_OK,
    THREAD_STATUS_SCHEDULED_REENTRY, THREAD_STATUS_YIELD, ThreadState,
};
use crate::thread::{LUA_COERR, LUA_COFIN, LUA_COSUS, Thread};
use crate::types::LUA_TFUNCTION;

static CO_STATUS_NAMES: [&str; 5] = ["running", "suspended", "normal", "dead", "dead"];

static CO_FUNCS: [NativeFunction; 7] = [
    NativeFunction {
        name: "create",
        function: coroutine_create,
    },
    NativeFunction {
        name: "running",
        function: coroutine_running,
    },
    NativeFunction {
        name: "status",
        function: coroutine_status,
    },
    NativeFunction {
        name: "wrap",
        function: coroutine_wrap,
    },
    NativeFunction {
        name: "yield",
        function: coroutine_yield,
    },
    NativeFunction {
        name: "isyieldable",
        function: coroutine_is_yieldable,
    },
    NativeFunction {
        name: "close",
        function: coroutine_close,
    },
];

enum CoroutineResume {
    Values(usize),
    Error,
}

/// `costatus`
fn coroutine_status(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let co = thread.to_thread(1);
        ctx.arg(1).expected(co.is_some(), "thread")?;
        let co = co.unwrap_unchecked();

        thread.push_string(CO_STATUS_NAMES[thread.co_status(&co) as usize])?;
        Ok(1)
    }
}

/// `auxresume`
unsafe fn aux_resume(thread: &Thread, co: &Thread, narg: i32) -> VmResult<CoroutineResume> {
    unsafe {
        if co.as_ptr().as_ref().unwrap_unchecked().status != THREAD_STATUS_YIELD {
            let status = thread.co_status(co);
            if status != LUA_COSUS {
                let _ = crate::push_fstring!(
                    thread,
                    "cannot resume %s coroutine",
                    CO_STATUS_NAMES[status as usize]
                );
                return Ok(CoroutineResume::Error);
            }
        }

        if narg != 0 {
            if co.check_stack(narg) == 0 {
                return crate::error!(thread, "too many arguments to resume").map_err(Into::into);
            }
            thread.x_move(co, narg)?;
        } else if co.stack_top().offset_from(co.stack_base()) as i32
            > crate::thread::LUAI_MAX_C_STACK
        {
            return crate::error!(thread, "too many arguments to resume").map_err(Into::into);
        }

        co.as_ptr().as_mut().unwrap_unchecked().single_step =
            thread.as_ptr().as_ref().unwrap_unchecked().single_step;

        match co.resume(Some(thread), narg) {
            Ok(()) | Err(VmExit::Control(VmControl::Yield)) => {
                debug_assert_ne!(
                    co.as_ptr().as_ref().unwrap_unchecked().status,
                    THREAD_STATUS_SCHEDULED_REENTRY
                );

                let nres = co.stack_top().offset_from(co.stack_base()) as i32;
                if nres != 0 {
                    if nres + 1 > crate::thread::LUA_MIN_STACK as i32
                        && thread.check_stack(nres + 1) == 0
                    {
                        return crate::error!(thread, "too many results to resume")
                            .map_err(Into::into);
                    }
                    co.x_move(thread, nres)?;
                }
                Ok(CoroutineResume::Values(nres as usize))
            }
            Err(exit) => {
                let VmExit::Error(_) = exit else {
                    return Err(exit);
                };
                co.x_move(thread, 1)?;
                Ok(CoroutineResume::Error)
            }
        }
    }
}

/// `interruptThread`
unsafe fn interrupt_thread(thread: &Thread, co: &Thread) -> NativeCallResult {
    unsafe {
        let global = thread.global();
        if let Some(hook) = global.debug_interrupt_callback() {
            thread.call_hook(
                |thread, debug| hook(thread, debug).map_err(Into::into),
                co.as_ptr().cast(),
            )?;
        }

        thread.break_current()
    }
}

/// `auxresumecont`
unsafe fn aux_resume_cont(thread: &Thread, co: &Thread) -> VmResult<CoroutineResume> {
    unsafe {
        let status = co.as_ptr().as_ref().unwrap_unchecked().status;
        if matches!(status, x if x == THREAD_STATUS_OK || x == THREAD_STATUS_YIELD) {
            let nres = co.stack_top().offset_from(co.stack_base()) as i32;
            if thread.check_stack(nres + 1) == 0 {
                return crate::error!(thread, "too many results to resume").map_err(Into::into);
            }
            co.x_move(thread, nres)?;
            Ok(CoroutineResume::Values(nres as usize))
        } else {
            thread.raw_check_stack(2)?;
            co.x_move(thread, 1)?;
            Ok(CoroutineResume::Error)
        }
    }
}

/// `coresumefinish`
unsafe fn coroutine_resume_finish(thread: &Thread, result: CoroutineResume) -> NativeCallResult {
    unsafe {
        match result {
            CoroutineResume::Values(count) => {
                let count_i32 = i32::try_from(count).expect("coroutine result count fits in i32");
                thread.push_boolean(1)?;
                thread.insert(-(count_i32 + 1));
                Ok(count + 1)
            }
            CoroutineResume::Error => {
                thread.push_boolean(0)?;
                thread.insert(-2);
                Ok(2)
            }
        }
    }
}

/// `auxwrapfinish`
unsafe fn aux_wrap_finish(thread: &Thread, result: CoroutineResume) -> NativeCallResult {
    unsafe {
        match result {
            CoroutineResume::Values(count) => Ok(count),
            CoroutineResume::Error => {
                if thread.is_string(-1) != 0 {
                    thread.push_where(1)?;
                    thread.insert(-2);
                    thread.concat(2)?;
                }
                thread.error().map_err(Into::into)
            }
        }
    }
}

/// `coresumey`
fn coroutine_resume_yieldable(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let co = thread.to_thread(1);
        ctx.arg(1).expected(co.is_some(), "thread")?;
        let co = co.unwrap_unchecked();
        let narg = thread.get_top() - 1;
        let result = match aux_resume(thread, &co, narg) {
            Ok(result) => result,
            Err(VmExit::Control(VmControl::Break)) => return interrupt_thread(thread, &co),
            Err(exit) => return Err(exit),
        };

        coroutine_resume_finish(thread, result)
    }
}

/// `coresumecont`
fn coroutine_resume_cont(ctx: NativeCallContext, _status: i32) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let co = thread.to_thread(1);
        ctx.arg(1).expected(co.is_some(), "thread")?;
        let co = co.unwrap_unchecked();

        if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_BREAK {
            return interrupt_thread(thread, &co);
        }

        let result = aux_resume_cont(thread, &co)?;
        coroutine_resume_finish(thread, result)
    }
}

/// `auxwrapy`
fn coroutine_wrap_yieldable(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let co = thread
            .to_thread(crate::thread::upvalue_index(1))
            .unwrap_unchecked();
        let result = match aux_resume(thread, &co, thread.get_top()) {
            Ok(result) => result,
            Err(VmExit::Control(VmControl::Break)) => return interrupt_thread(thread, &co),
            Err(exit) => return Err(exit),
        };

        aux_wrap_finish(thread, result)
    }
}

/// `auxwrapcont`
fn coroutine_wrap_cont(ctx: NativeCallContext, _status: i32) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let co = thread
            .to_thread(crate::thread::upvalue_index(1))
            .unwrap_unchecked();

        if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_BREAK {
            return interrupt_thread(thread, &co);
        }

        let result = aux_resume_cont(thread, &co)?;
        aux_wrap_finish(thread, result)
    }
}

/// `cocreate`
fn coroutine_create(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        thread.check_type(1, LUA_TFUNCTION)?;
        let new_thread = thread.new_thread()?;
        thread.x_push(&new_thread, 1)?;
        Ok(1)
    }
}

/// `cowrap`
fn coroutine_wrap(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    coroutine_create(NativeCallContext::new(thread))?;
    unsafe {
        thread.push_native_closure_k(
            coroutine_wrap_yieldable,
            None,
            1,
            Some(coroutine_wrap_cont),
        )?
    };
    Ok(1)
}

/// `coyield`
fn coroutine_yield(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    let nres = unsafe { thread.get_top() };
    unsafe { thread.yield_current(nres) }
}

/// `corunning`
fn coroutine_running(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    if unsafe { thread.push_thread()? } != 0 {
        unsafe { thread.push_nil()? };
    }
    Ok(1)
}

/// `coyieldable`
fn coroutine_is_yieldable(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe { thread.push_boolean(thread.is_yieldable())? };
    Ok(1)
}

/// `coclose`
fn coroutine_close(ctx: NativeCallContext) -> NativeCallResult {
    let thread = ctx.raw_thread();
    unsafe {
        let co = thread.to_thread(1);
        ctx.arg(1).expected(co.is_some(), "thread")?;
        let co = co.unwrap_unchecked();

        let status = thread.co_status(&co);
        if status != LUA_COFIN && status != LUA_COERR && status != LUA_COSUS {
            return crate::error!(
                thread,
                "cannot close %s coroutine",
                CO_STATUS_NAMES[status as usize]
            )
            .map_err(Into::into);
        }

        if matches!(
            co.as_ptr().as_ref().unwrap_unchecked().status,
            x if x == THREAD_STATUS_OK || x == THREAD_STATUS_YIELD
        ) {
            thread.push_boolean(1)?;
            co.reset()?;
            Ok(1)
        } else {
            thread.push_boolean(0)?;

            if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_ERR_MEM {
                thread.push_string(crate::state::LUA_MEMERRMSG)?;
            } else if co.as_ptr().as_ref().unwrap_unchecked().status == THREAD_STATUS_ERR_ERR {
                thread.push_string(crate::state::LUA_ERRERRMSG)?;
            } else if co.get_top() != 0 {
                co.x_move(thread, 1)?;
            }

            co.reset()?;
            Ok(2)
        }
    }
}

impl Thread {
    /// `luaopen_coroutine`
    pub unsafe fn open_coroutine(&self) -> NativeCallResult {
        unsafe { self.register(Some(super::LUA_COLIB_NAME), &CO_FUNCS[..])? };
        unsafe {
            self.push_native_closure_k(
                coroutine_resume_yieldable,
                Some("resume"),
                0,
                Some(coroutine_resume_cont),
            )?;
            self.raw_set_field(-2, "resume")?;
        }
        Ok(1)
    }
}