From 743d1075bf6ba3131116d4e177922368be534d8d Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Thu, 1 Feb 2024 14:07:03 +0100
Subject: [PATCH] Add exit lua function to functions struct

---
 examples/exit_code.rs | 42 ++++++------------------------------------
 lib/functions.rs      | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 36 deletions(-)

diff --git a/examples/exit_code.rs b/examples/exit_code.rs
index 7d9be1d..2e9b664 100644
--- a/examples/exit_code.rs
+++ b/examples/exit_code.rs
@@ -1,54 +1,24 @@
 #![allow(clippy::missing_errors_doc)]
 #![allow(clippy::missing_panics_doc)]
 
-use std::process::ExitCode;
-
 use async_io::block_on;
 
 use mlua::prelude::*;
-use mlua_luau_runtime::{LuaRuntimeExt, Runtime};
+use mlua_luau_runtime::{Functions, Runtime};
 
 const MAIN_SCRIPT: &str = include_str!("./lua/exit_code.luau");
 
-const EXIT_IMPL_LUA: &str = r"
-exit(...)
-yield()
-";
-
 pub fn main() -> LuaResult<()> {
     tracing_subscriber::fmt::init();
 
     // Set up persistent Lua environment
     let lua = Lua::new();
-
-    // Note that our exit function is partially implemented in Lua
-    // because we need to also yield the thread that called it, this
-    // is not possible to do in Rust because of crossing C-call boundary
-    let exit_fn_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_fn = lua
-        .load(EXIT_IMPL_LUA)
-        .set_environment(exit_fn_env)
-        .into_function()?;
-    lua.globals().set("exit", exit_fn)?;
-
-    // Load the main script into a runtime
     let rt = Runtime::new(&lua);
+    let fns = Functions::new(&lua)?;
+
+    lua.globals().set("exit", fns.exit)?;
+
+    // Load the main script into the runtime
     let main = lua.load(MAIN_SCRIPT);
     rt.push_thread_front(main, ())?;
 
diff --git a/lib/functions.rs b/lib/functions.rs
index e52160f..876d76d 100644
--- a/lib/functions.rs
+++ b/lib/functions.rs
@@ -1,6 +1,8 @@
 #![allow(unused_imports)]
 #![allow(clippy::module_name_repetitions)]
 
+use std::process::ExitCode;
+
 use mlua::prelude::*;
 
 use crate::{
@@ -9,6 +11,7 @@ use crate::{
     result_map::ThreadResultMap,
     runtime::Runtime,
     thread_id::ThreadId,
+    traits::LuaRuntimeExt,
     util::{is_poll_pending, LuaThreadOrFunction, ThreadResult},
 };
 
@@ -18,6 +21,11 @@ Lua state does not have runtime metadata attached!\
 \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`].
 */
@@ -38,6 +46,12 @@ pub struct Functions<'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> {
@@ -128,10 +142,33 @@ impl<'lua> Functions<'lua> {
             }
         })?;
 
+        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,
         })
     }
 }