From d7404679c7c89a7f8c956e04c821f12278e9ecd3 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Sun, 20 Aug 2023 21:15:22 -0500
Subject: [PATCH] More robust task.spawn implementation

---
 src/lune/builtins/task/mod.rs | 61 +++++++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 14 deletions(-)

diff --git a/src/lune/builtins/task/mod.rs b/src/lune/builtins/task/mod.rs
index 2bba0f8..71e74f4 100644
--- a/src/lune/builtins/task/mod.rs
+++ b/src/lune/builtins/task/mod.rs
@@ -9,12 +9,58 @@ use crate::lune::{scheduler::Scheduler, util::TableBuilder};
 mod tof;
 use tof::LuaThreadOrFunction;
 
+/*
+    The spawn function needs special treatment,
+    we need to yield right away to allow the
+    spawned task to run until first yield
+
+    1. Schedule this current thread at the front
+    2. Schedule given thread/function at the front,
+       the previous schedule now comes right after
+    3. Give control over to the scheduler, which will
+       resume the above tasks in order when its ready
+*/
+const SPAWN_IMPL_LUA: &str = r#"
+push(currentThread())
+local thread = push(...)
+yield()
+return thread
+"#;
+
 pub fn create(lua: &'static Lua) -> LuaResult<LuaTable<'_>> {
+    let coroutine_running = lua
+        .globals()
+        .get::<_, LuaTable>("coroutine")?
+        .get::<_, LuaFunction>("running")?;
+    let coroutine_yield = lua
+        .globals()
+        .get::<_, LuaTable>("coroutine")?
+        .get::<_, LuaFunction>("yield")?;
+    let push_front =
+        lua.create_function(|lua, (tof, args): (LuaThreadOrFunction, LuaMultiValue)| {
+            let thread = tof.into_thread(lua)?;
+            let sched = lua
+                .app_data_ref::<&Scheduler>()
+                .expect("Lua struct is missing scheduler");
+            sched.push_front(thread.clone(), args)?;
+            Ok(thread)
+        })?;
+    let task_spawn_env = TableBuilder::new(lua)?
+        .with_value("currentThread", coroutine_running)?
+        .with_value("yield", coroutine_yield)?
+        .with_value("push", push_front)?
+        .build_readonly()?;
+    let task_spawn = lua
+        .load(SPAWN_IMPL_LUA)
+        .set_name("task.spawn")
+        .set_environment(task_spawn_env)
+        .into_function()?;
+
     TableBuilder::new(lua)?
         .with_function("cancel", task_cancel)?
         .with_function("defer", task_defer)?
         .with_function("delay", task_delay)?
-        .with_function("spawn", task_spawn)?
+        .with_value("spawn", task_spawn)?
         .with_async_function("wait", task_wait)?
         .build_readonly()
 }
@@ -69,19 +115,6 @@ where
     Ok(thread)
 }
 
-fn task_spawn<'lua>(
-    lua: &'lua Lua,
-    (tof, args): (LuaThreadOrFunction<'lua>, LuaMultiValue<'_>),
-) -> LuaResult<LuaThread<'lua>> {
-    let thread = tof.into_thread(lua)?;
-    let resume = lua
-        .globals()
-        .get::<_, LuaTable>("coroutine")?
-        .get::<_, LuaFunction>("resume")?;
-    resume.call((thread.clone(), args))?;
-    Ok(thread)
-}
-
 async fn task_wait(_: &Lua, secs: Option<f64>) -> LuaResult<f64> {
     let duration = Duration::from_secs_f64(secs.unwrap_or_default());