From 7c8af9730f9b036f636bdf730393a67b47afcde7 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Thu, 17 Aug 2023 00:23:20 -0500
Subject: [PATCH] Some work on async functions for new scheduler

---
 Cargo.lock                       | 40 ++++++++++++++++++++++++++++
 Cargo.toml                       |  1 +
 src/lune/scheduler/impl_async.rs | 29 +++++++++++++++++---
 src/lune/scheduler/traits.rs     | 45 +++++++++++++++++++++++++++++++-
 4 files changed, 111 insertions(+), 4 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c7e4220..b26b0f2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1140,6 +1140,7 @@ dependencies = [
  "bstr",
  "erased-serde",
  "mlua-sys",
+ "mlua_derive",
  "num-traits",
  "once_cell",
  "rustc-hash",
@@ -1159,6 +1160,21 @@ dependencies = [
  "pkg-config",
 ]
 
+[[package]]
+name = "mlua_derive"
+version = "0.9.0-rc.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "4b900bf5619cdf327722cb39424409856a46f3a516628117a896cdbeecc0a9b1"
+dependencies = [
+ "itertools",
+ "once_cell",
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "regex",
+ "syn 2.0.28",
+]
+
 [[package]]
 name = "nibble_vec"
 version = "0.1.0"
@@ -1331,6 +1347,30 @@ version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"
 
+[[package]]
+name = "proc-macro-error"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c"
+dependencies = [
+ "proc-macro-error-attr",
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+ "version_check",
+]
+
+[[package]]
+name = "proc-macro-error-attr"
+version = "1.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "version_check",
+]
+
 [[package]]
 name = "proc-macro-hack"
 version = "0.5.20+deprecated"
diff --git a/Cargo.toml b/Cargo.toml
index 0bc1d5c..deabb10 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -78,6 +78,7 @@ urlencoding = "2.1"
 ### RUNTIME
 
 mlua = { version = "0.9.0-beta.3", features = [
+    "macros",
     "luau",
     "luau-jit",
     "serialize",
diff --git a/src/lune/scheduler/impl_async.rs b/src/lune/scheduler/impl_async.rs
index 4b2d279..8ded33b 100644
--- a/src/lune/scheduler/impl_async.rs
+++ b/src/lune/scheduler/impl_async.rs
@@ -1,18 +1,41 @@
 use futures_util::Future;
+use mlua::prelude::*;
 
-use super::SchedulerImpl;
+use super::{traits::IntoLuaThread, SchedulerImpl};
 
-impl SchedulerImpl {
+impl<'lua> SchedulerImpl {
     /**
         Schedules a plain future to run whenever the scheduler is available.
     */
     pub fn schedule_future<F>(&self, fut: F)
     where
-        F: Future<Output = ()> + 'static,
+        F: 'static + Future<Output = ()>,
     {
         self.futures
             .try_lock()
             .expect("Failed to lock futures queue")
             .push(Box::pin(fut))
     }
+    /**
+        Schedules the given `thread` to run when the given `fut` completes.
+    */
+    pub fn schedule_thread<T, R, F>(&'lua self, thread: T, fut: F) -> LuaResult<()>
+    where
+        T: IntoLuaThread<'lua>,
+        R: IntoLuaMulti<'lua>,
+        F: 'static + Future<Output = LuaResult<R>>,
+    {
+        let thread = thread.into_lua_thread(&self.lua)?;
+
+        let fut = async move {
+            let rets = fut.await.expect("Failed to receive result");
+            self.push_back(thread, rets)
+                .expect("Failed to schedule future thread");
+        };
+
+        // TODO: Lifetime issues
+        // self.schedule_future(fut);
+
+        Ok(())
+    }
 }
diff --git a/src/lune/scheduler/traits.rs b/src/lune/scheduler/traits.rs
index 5a05d5d..3901458 100644
--- a/src/lune/scheduler/traits.rs
+++ b/src/lune/scheduler/traits.rs
@@ -1,4 +1,5 @@
-use mlua::prelude::*;
+use futures_util::Future;
+use mlua::{chunk, prelude::*};
 
 use super::Scheduler;
 
@@ -15,6 +16,20 @@ pub trait LuaSchedulerExt {
         not be dropped either because of the strong reference.
     */
     fn scheduler(&self) -> Scheduler;
+
+    /**
+        Creates a function callable from Lua that runs an async
+        closure and returns the results of it to the call site.
+    */
+    fn create_async_function<'lua, A, R, F, FR>(
+        &'lua self,
+        func: F,
+    ) -> LuaResult<LuaFunction<'lua>>
+    where
+        A: FromLuaMulti<'lua>,
+        R: IntoLuaMulti<'lua>,
+        F: 'static + Fn(&'lua Lua, A) -> FR,
+        FR: 'static + Future<Output = LuaResult<R>>;
 }
 
 impl LuaSchedulerExt for Lua {
@@ -23,6 +38,34 @@ impl LuaSchedulerExt for Lua {
             .expect("Lua struct is missing scheduler")
             .clone()
     }
+
+    fn create_async_function<'lua, A, R, F, FR>(&'lua self, func: F) -> LuaResult<LuaFunction<'lua>>
+    where
+        A: FromLuaMulti<'lua>,
+        R: IntoLuaMulti<'lua>,
+        F: 'static + Fn(&'lua Lua, A) -> FR,
+        FR: 'static + Future<Output = LuaResult<R>>,
+    {
+        let async_yield = self
+            .globals()
+            .get::<_, LuaTable>("coroutine")?
+            .get::<_, LuaFunction>("yield")?;
+        let async_schedule = self.create_function(move |lua: &Lua, args: A| {
+            let thread = lua.current_thread();
+            let future = func(lua, args);
+            // TODO: Add to scheduler
+            Ok(())
+        })?;
+
+        let async_func = self
+            .load(chunk!({
+                $async_schedule(...)
+                return $async_yield()
+            }))
+            .set_name("async")
+            .into_function()?;
+        Ok(async_func)
+    }
 }
 
 /**