From 5820858147415dce90587d20288064dc501a912e Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Wed, 31 Jan 2024 19:40:11 +0100
Subject: [PATCH] Implement thread id struct and some more utils

---
 lib/handle.rs    | 38 +++++++++++++-------------------------
 lib/lib.rs       |  2 ++
 lib/thread_id.rs | 32 ++++++++++++++++++++++++++++++++
 lib/util.rs      | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 80 insertions(+), 25 deletions(-)
 create mode 100644 lib/thread_id.rs

diff --git a/lib/handle.rs b/lib/handle.rs
index d23671a..56b3458 100644
--- a/lib/handle.rs
+++ b/lib/handle.rs
@@ -14,7 +14,7 @@ use crate::{
     runtime::Runtime,
     status::Status,
     traits::IntoLuaThread,
-    util::{run_until_yield, ThreadWithArgs},
+    util::{run_until_yield, ThreadResult, ThreadWithArgs},
 };
 
 /**
@@ -28,7 +28,7 @@ use crate::{
 #[derive(Debug, Clone)]
 pub struct Handle {
     thread: Rc<RefCell<Option<ThreadWithArgs>>>,
-    result: Rc<RefCell<Option<(bool, LuaRegistryKey)>>>,
+    result: Rc<RefCell<Option<ThreadResult>>>,
     status: Rc<Cell<bool>>,
     event: Rc<Event>,
 }
@@ -69,24 +69,14 @@ impl Handle {
             .into_inner(lua)
     }
 
-    fn set<'lua>(
-        &self,
-        lua: &'lua Lua,
-        result: &LuaResult<LuaMultiValue<'lua>>,
-        is_final: bool,
-    ) -> LuaResult<()> {
-        self.result.borrow_mut().replace((
-            result.is_ok(),
-            match &result {
-                Ok(v) => lua.create_registry_value(v.clone().into_vec())?,
-                Err(e) => lua.create_registry_value(e.clone())?,
-            },
-        ));
+    fn set<'lua>(&self, lua: &'lua Lua, result: &LuaResult<LuaMultiValue<'lua>>, is_final: bool) {
+        self.result
+            .borrow_mut()
+            .replace(ThreadResult::new(result.clone(), lua));
         self.status.replace(is_final);
         if is_final {
             self.event.notify(usize::MAX);
         }
-        Ok(())
     }
 
     /**
@@ -97,17 +87,15 @@ impl Handle {
         - [`Status::NotStarted`]: returns `None`.
         - [`Status::Running`]: may return `Some(Ok(v))` or `Some(Err(e))`, but it is not guaranteed.
         - [`Status::Completed`]: returns `Some(Ok(v))` or `Some(Err(e))`.
+
+        Note that this method also takes the value out of the handle, so it may only be called once.
+
+        Any subsequent calls after this method returns `Some` will return `None`.
     */
     #[must_use]
     pub fn result<'lua>(&self, lua: &'lua Lua) -> Option<LuaResult<LuaMultiValue<'lua>>> {
-        let res = self.result.borrow();
-        let (is_ok, key) = res.as_ref()?;
-        Some(if *is_ok {
-            let v = lua.registry_value(key).unwrap();
-            Ok(LuaMultiValue::from_vec(v))
-        } else {
-            Err(lua.registry_value(key).unwrap())
-        })
+        let mut res = self.result.borrow_mut();
+        res.take().map(|r| r.value(lua))
     }
 
     /**
@@ -135,7 +123,7 @@ impl LuaUserData for Handle {
             let (thread, args) = this.take(lua);
             let result = run_until_yield(thread.clone(), args).await;
             let is_final = thread.status() != LuaThreadStatus::Resumable;
-            this.set(lua, &result, is_final)?;
+            this.set(lua, &result, is_final);
             result
         });
     }
diff --git a/lib/lib.rs b/lib/lib.rs
index ccc41c9..d17a15e 100644
--- a/lib/lib.rs
+++ b/lib/lib.rs
@@ -4,6 +4,7 @@ mod handle;
 mod queue;
 mod runtime;
 mod status;
+mod thread_id;
 mod traits;
 mod util;
 
@@ -11,4 +12,5 @@ pub use functions::Functions;
 pub use handle::Handle;
 pub use runtime::Runtime;
 pub use status::Status;
+pub use thread_id::ThreadId;
 pub use traits::{IntoLuaThread, LuaRuntimeExt};
diff --git a/lib/thread_id.rs b/lib/thread_id.rs
new file mode 100644
index 0000000..b5afed4
--- /dev/null
+++ b/lib/thread_id.rs
@@ -0,0 +1,32 @@
+use std::hash::{Hash, Hasher};
+
+use mlua::prelude::*;
+
+/**
+    Opaque and unique ID representing a [`LuaThread`].
+
+    Typically used for associating metadata with a thread in a structure such as a `HashMap<ThreadId, ...>`.
+
+    Note that holding a `ThreadId` does not prevent the thread from being garbage collected.
+    The actual thread may or may not still exist and be active at any given point in time.
+*/
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ThreadId(usize);
+
+impl From<LuaThread<'_>> for ThreadId {
+    fn from(thread: LuaThread) -> Self {
+        Self(LuaValue::Thread(thread).to_pointer() as usize)
+    }
+}
+
+impl From<&LuaThread<'_>> for ThreadId {
+    fn from(thread: &LuaThread) -> Self {
+        Self(LuaValue::Thread(thread.clone()).to_pointer() as usize)
+    }
+}
+
+impl Hash for ThreadId {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.0.hash(state);
+    }
+}
diff --git a/lib/util.rs b/lib/util.rs
index 5001901..933886b 100644
--- a/lib/util.rs
+++ b/lib/util.rs
@@ -21,6 +21,39 @@ pub(crate) async fn run_until_yield<'lua>(
     stream.next().await.unwrap()
 }
 
+/**
+    Representation of a [`LuaResult`] with an associated [`LuaMultiValue`] currently stored in the Lua registry.
+*/
+#[derive(Debug)]
+pub(crate) struct ThreadResult {
+    inner: LuaResult<LuaRegistryKey>,
+}
+
+impl ThreadResult {
+    pub fn new(result: LuaResult<LuaMultiValue>, lua: &Lua) -> Self {
+        Self {
+            inner: match result {
+                Ok(v) => Ok({
+                    let vec = v.into_vec();
+                    lua.create_registry_value(vec).expect("out of memory")
+                }),
+                Err(e) => Err(e),
+            },
+        }
+    }
+
+    pub fn value(self, lua: &Lua) -> LuaResult<LuaMultiValue> {
+        match self.inner {
+            Ok(key) => {
+                let vec = lua.registry_value(&key).unwrap();
+                lua.remove_registry_value(key).unwrap();
+                Ok(LuaMultiValue::from_vec(vec))
+            }
+            Err(e) => Err(e.clone()),
+        }
+    }
+}
+
 /**
     Representation of a [`LuaThread`] with its associated arguments currently stored in the Lua registry.
 */