From 0c657ae78e95914afcbe799d80a7636ca91614a2 Mon Sep 17 00:00:00 2001
From: Discord9 <discord9@163.com>
Date: Mon, 29 Jan 2024 11:26:34 +0800
Subject: [PATCH] feat(hydroflow): add `StateHandleErased` (no type param) for
 storage in data structures, fix #1048

---
 hydroflow/src/scheduled/mod.rs   |  2 +-
 hydroflow/src/scheduled/state.rs | 74 +++++++++++++++++++++++++++++++-
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/hydroflow/src/scheduled/mod.rs b/hydroflow/src/scheduled/mod.rs
index 17417788dddf..e414fdfab324 100644
--- a/hydroflow/src/scheduled/mod.rs
+++ b/hydroflow/src/scheduled/mod.rs
@@ -43,6 +43,6 @@ impl Display for HandoffId {
 
 /// A staten handle's ID. Invalid if used in a different [`graph::Hydroflow`]
 /// instance than the original that created it.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 #[repr(transparent)]
 pub struct StateId(pub(crate) usize);
diff --git a/hydroflow/src/scheduled/state.rs b/hydroflow/src/scheduled/state.rs
index 05c862cd06d1..f75076c2a377 100644
--- a/hydroflow/src/scheduled/state.rs
+++ b/hydroflow/src/scheduled/state.rs
@@ -1,13 +1,17 @@
 //! Module for [`StateHandle`], part of the "state API".
 
+use std::any::{Any, TypeId};
 use std::marker::PhantomData;
 
 use super::StateId;
 
 /// A handle into a particular [`Hydroflow`](super::graph::Hydroflow) instance, referring to data
 /// inserted by [`add_state`](super::graph::Hydroflow::add_state).
+///
+/// If you need to store state handles in a data structure see [`StateHandleErased`] which hides
+/// the generic type parameter.
 #[must_use]
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct StateHandle<T> {
     pub(crate) state_id: StateId,
     pub(crate) _phantom: PhantomData<*mut T>,
@@ -18,3 +22,71 @@ impl<T> Clone for StateHandle<T> {
         *self
     }
 }
+
+/// A state handle with the generic type parameter erased, allowing it to be stored in omogenous
+/// data structures. The type is tracked internally as data via [`TypeId`].
+///
+/// Use [`StateHandleErased::from(state_handle)`](StateHandleErased::from) to create an instance
+/// from a typed [`StateHandle<T>`].
+///
+/// Use [`StateHandle::<T>::try_from()`](StateHandle::try_from) to convert the `StateHandleErased`
+/// back into a `StateHandle<T>` of the given type `T`. If `T` is the wrong type then the original
+/// `StateHandleErased` will be returned as the `Err`.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
+pub struct StateHandleErased {
+    state_id: StateId,
+    type_id: TypeId,
+}
+
+/// See [`StateHandleErased`].
+impl<T> TryFrom<StateHandleErased> for StateHandle<T>
+where
+    T: Any,
+{
+    type Error = StateHandleErased;
+
+    fn try_from(value: StateHandleErased) -> Result<Self, Self::Error> {
+        if TypeId::of::<T>() == value.type_id {
+            Ok(Self {
+                state_id: value.state_id,
+                _phantom: PhantomData,
+            })
+        } else {
+            Err(value)
+        }
+    }
+}
+/// See [`StateHandleErased`].
+impl<T> From<StateHandle<T>> for StateHandleErased
+where
+    T: Any,
+{
+    fn from(value: StateHandle<T>) -> Self {
+        Self {
+            state_id: value.state_id,
+            type_id: TypeId::of::<T>(),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use std::marker::PhantomData;
+
+    use super::*;
+    use crate::scheduled::StateId;
+
+    #[test]
+    fn test_erasure() {
+        let handle = StateHandle::<String> {
+            state_id: StateId(0),
+            _phantom: PhantomData,
+        };
+        let handle_erased = StateHandleErased::from(handle);
+        let handle_good = StateHandle::<String>::try_from(handle_erased);
+        let handle_bad = StateHandle::<&'static str>::try_from(handle_erased);
+
+        assert_eq!(Ok(handle), handle_good);
+        assert_eq!(Err(handle_erased), handle_bad);
+    }
+}