mlua-luau-scheduler/lib/functions.rs

#![allow(unused_imports)]
#![allow(clippy::module_name_repetitions)]

use std::process::ExitCode;

use mlua::prelude::*;

use crate::{
    error_callback::ThreadErrorCallback,
    queue::{DeferredThreadQueue, SpawnedThreadQueue},
    result_map::ThreadResultMap,
    runtime::Runtime,
    thread_id::ThreadId,
    traits::LuaRuntimeExt,
    util::{is_poll_pending, LuaThreadOrFunction, ThreadResult},
};

const ERR_METADATA_NOT_ATTACHED: &str = "\
Lua state does not have runtime metadata attached!\
\nThis is most likely caused by creating functions outside of a runtime.\
\nRuntime functions must always be created from within an active runtime.\
";

const EXIT_IMPL_LUA: &str = r"
exit(...)
yield()
";

/**
    A collection of lua functions that may be called to interact with a [`Runtime`].
*/
pub struct Functions<'lua> {
    /**
        Resumes a function / thread once instantly, and runs until first yield.

        Spawns onto the runtime queue if not completed.
    */
    pub spawn: LuaFunction<'lua>,
    /**
        Defers a function / thread onto the runtime queue.

        Does not resume instantly, only adds to the queue.
    */
    pub defer: LuaFunction<'lua>,
    /**
        Cancels a function / thread, removing it from the queue.
    */
    pub cancel: LuaFunction<'lua>,
    /**
        Exits the runtime, stopping all other threads and closing the runtime.

        Yields the calling thread to ensure that it does not continue.
    */
    pub exit: LuaFunction<'lua>,
}

impl<'lua> Functions<'lua> {
    /**
        Creates a new collection of Lua functions that may be called to interact with a [`Runtime`].

        # Errors

        Errors when out of memory, or if default Lua globals are missing.

        # Panics

        Panics when the given [`Lua`] instance does not have an attached [`Runtime`].
    */
    pub fn new(lua: &'lua Lua) -> LuaResult<Self> {
        let spawn_queue = lua
            .app_data_ref::<SpawnedThreadQueue>()
            .expect(ERR_METADATA_NOT_ATTACHED)
            .clone();
        let defer_queue = lua
            .app_data_ref::<DeferredThreadQueue>()
            .expect(ERR_METADATA_NOT_ATTACHED)
            .clone();
        let error_callback = lua
            .app_data_ref::<ThreadErrorCallback>()
            .expect(ERR_METADATA_NOT_ATTACHED)
            .clone();
        let result_map = lua
            .app_data_ref::<ThreadResultMap>()
            .expect(ERR_METADATA_NOT_ATTACHED)
            .clone();

        let spawn_map = result_map.clone();
        let spawn = lua.create_function(
            move |lua, (tof, args): (LuaThreadOrFunction, LuaMultiValue)| {
                let thread = tof.into_thread(lua)?;
                if thread.status() == LuaThreadStatus::Resumable {
                    // NOTE: We need to resume the thread once instantly for correct behavior,
                    // and only if we get the pending value back we can spawn to async executor
                    match thread.resume::<_, LuaMultiValue>(args.clone()) {
                        Ok(v) => {
                            if v.get(0).map(is_poll_pending).unwrap_or_default() {
                                spawn_queue.push_item(lua, &thread, args)?;
                            } else {
                                // Not pending, store the value
                                let id = ThreadId::from(&thread);
                                if spawn_map.is_tracked(id) {
                                    let res = ThreadResult::new(Ok(v), lua);
                                    spawn_map.insert(id, res);
                                }
                            }
                        }
                        Err(e) => {
                            error_callback.call(&e);
                            // Not pending, store the error
                            let id = ThreadId::from(&thread);
                            if spawn_map.is_tracked(id) {
                                let res = ThreadResult::new(Err(e), lua);
                                spawn_map.insert(id, res);
                            }
                        }
                    };
                }
                Ok(thread)
            },
        )?;

        let defer = lua.create_function(
            move |lua, (tof, args): (LuaThreadOrFunction, LuaMultiValue)| {
                let thread = tof.into_thread(lua)?;
                if thread.status() == LuaThreadStatus::Resumable {
                    defer_queue.push_item(lua, &thread, args)?;
                }
                Ok(thread)
            },
        )?;

        let close = lua
            .globals()
            .get::<_, LuaTable>("coroutine")?
            .get::<_, LuaFunction>("close")?;
        let close_key = lua.create_registry_value(close)?;
        let cancel = lua.create_function(move |lua, thread: LuaThread| {
            let close: LuaFunction = lua.registry_value(&close_key)?;
            match close.call(thread) {
                Err(LuaError::CoroutineInactive) | Ok(()) => Ok(()),
                Err(e) => Err(e),
            }
        })?;

        let exit_env = lua.create_table_from(vec![
            (
                "exit",
                lua.create_function(|lua, code: Option<u8>| {
                    let code = code.map(ExitCode::from).unwrap_or_default();
                    lua.set_exit_code(code);
                    Ok(())
                })?,
            ),
            (
                "yield",
                lua.globals()
                    .get::<_, LuaTable>("coroutine")?
                    .get::<_, LuaFunction>("yield")?,
            ),
        ])?;
        let exit = lua
            .load(EXIT_IMPL_LUA)
            .set_name("=__runtime_exit")
            .set_environment(exit_env)
            .into_function()?;

        Ok(Self {
            spawn,
            defer,
            cancel,
            exit,
        })
    }
}