Add extend method for UserDataMap

DelSkayn · Aug 3, 2024 · 2c2c220 · 2c2c220
1 parent 4dabdc2
commit 2c2c220
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 3 deletions.
diff --git a/core/src/context/ctx.rs b/core/src/context/ctx.rs
@@ -15,7 +15,7 @@ use crate::{
     cstr,
     markers::Invariant,
     qjs,
-    runtime::{opaque::Opaque, UserData, UserDataError, UserDataGuard},
+    runtime::{opaque::Opaque, UserData, UserDataError, UserDataGuard, UserDataMap},
     Atom, Error, FromJs, Function, IntoJs, Object, Promise, Result, String, Value,
 };
 
@@ -468,6 +468,13 @@ impl<'js> Ctx<'js> {
         unsafe { self.get_opaque().insert_userdata(data) }
     }
 
+    /// Store multiple types in the runtime which can be retrieved later with `Ctx::userdata`.
+    ///
+    /// Returns an error if the userdata is currently being accessed.
+    pub fn extend_userdata(&self, data: UserDataMap) -> StdResult<(), UserDataError<()>> {
+        unsafe { self.get_opaque().extend_userdata(data) }
+    }
+
     /// Remove the userdata of the given type from the userdata storage.
     ///
     /// Returns Err(()) if the userdata is currently being accessed and removing isn't possible.

diff --git a/core/src/runtime.rs b/core/src/runtime.rs
@@ -13,7 +13,7 @@ pub(crate) mod schedular;
 mod spawner;
 
 pub use base::{Runtime, WeakRuntime};
-pub use userdata::{UserData, UserDataError, UserDataGuard};
+pub use userdata::{UserData, UserDataError, UserDataGuard, UserDataMap};
 
 #[cfg(feature = "futures")]
 pub(crate) use r#async::InnerRuntime;

diff --git a/core/src/runtime/opaque.rs b/core/src/runtime/opaque.rs
@@ -94,6 +94,10 @@ impl<'js> Opaque<'js> {
         self.userdata.insert(data)
     }
 
+    pub fn extend_userdata(&self, data: UserDataMap) -> Result<(), UserDataError<()>> {
+        self.userdata.extend(data)
+    }
+
     pub fn remove_userdata<U>(&self) -> Result<Option<Box<U>>, UserDataError<()>>
     where
         U: UserData<'js>,

diff --git a/core/src/runtime/userdata.rs b/core/src/runtime/userdata.rs
@@ -152,7 +152,7 @@ impl Hasher for IdHasher {
 
 /// Typeid hashmap taken from axum.
 #[derive(Default)]
-pub(crate) struct UserDataMap {
+pub struct UserDataMap {
     map: UnsafeCell<HashMap<TypeId, Box<dyn Any>, BuildHasherDefault<IdHasher>>>,
     count: Cell<usize>,
 }
@@ -176,6 +176,14 @@ impl UserDataMap {
         Ok(r)
     }
 
+    pub fn extend(&self, other: Self) -> Result<(), UserDataError<()>> {
+        if self.count.get() > 0 || other.count.get() > 0 {
+            return Err(UserDataError(()));
+        }
+        unsafe { (*self.map.get()).extend((*other.map.get()).drain()) }
+        Ok(())
+    }
+
     pub fn remove<'js, U>(&self) -> Result<Option<Box<U>>, UserDataError<()>>
     where
         U: UserData<'js>,