Implement task handle awaiting

crates/lb/src/task.rs

@@ -1,5 +1,6 @@
 use crate::runtime::spawn;
 use luaffi::{cdef, metatype};
+use luajit::{LUA_MULTRET, Type};
 use std::{ffi::c_int, time::Duration};
 use tokio::{task::JoinHandle, time::sleep};
 
@@ -19,33 +20,67 @@ impl lb_tasklib {
     }
 
     pub extern "Lua" fn spawn(&self, f: function, ...) {
-        // pack the function and its arguments into a table and pass its ref to rust
+        // pack the function and its arguments into a table and pass its ref to rust.
+        //
+        // this table is used from rust-side to call the function with its args, and it's also
+        // reused to store its return values that the task handle can return when awaited. the ref
+        // is owned by the task handle and unref'ed when it's gc'ed.
         self.__spawn(__ref(__tpack(f, variadic!())))
     }
 
     extern "Lua-C" fn __spawn(&self, key: c_int) -> lb_task {
         let handle = spawn(async move |cx| {
-            // SAFETY: key is always unique, created by __ref above
+            // SAFETY: key is always unique, created by __ref above.
             let arg = unsafe { cx.new_ref_unchecked(key) };
             let mut s = cx.guard();
             s.resize(0);
-            s.push(arg);
-            let narg = s.unpack(1, 1, None) - 1; // unpack the table containing the function to call and its args
-            if let Err(err) = s.call_async(narg, 0).await {
-                drop(s);
-                cx.report_error(&err);
+            s.push(&arg);
+            let narg = s.unpack(1, 1, None) - 1; // unpack the function and its args from the table
+            debug_assert!(s.slot(2).type_of() == Type::Function);
+            match s.call_async(narg, LUA_MULTRET).await {
+                Ok(nret) => {
+                    s.pack(1, nret); // pack the return values back into the table
+                }
+                Err(err) => {
+                    drop(s);
+                    cx.report_error(&err);
+                }
             }
+            let _ = arg.into_raw(); // the original ref is owned by the task handle and unref'ed there
         });
 
-        lb_task { handle }
+        lb_task {
+            handle: Some(handle),
+            __ref: key,
+        }
     }
 }
 
 #[cdef]
 pub struct lb_task {
     #[opaque]
-    handle: JoinHandle<()>,
+    handle: Option<JoinHandle<()>>,
+    __ref: c_int,
 }
 
 #[metatype]
-impl lb_task {}
+impl lb_task {
+    pub extern "Lua" fn r#await(&self) -> many {
+        self.__await();
+        let ret = __registry[self.__ref];
+        __tunpack(ret, 1, ret.n)
+    }
+
+    async extern "Lua-C" fn __await(&mut self) {
+        if let Some(handle) = self.handle.take() {
+            handle
+                .await
+                .unwrap_or_else(|err| panic!("task handler panicked: {err}"));
+        }
+    }
+
+    #[gc]
+    extern "Lua" fn gc(&self) {
+        __unref(self.__ref);
+    }
+}