feat: add some func to TypeMap (#9)

cloudwego · Mar 1, 2023 · 38c8170 · 38c8170
1 parent ff933c1
commit 38c8170
Showing 1 changed file with 46 additions and 0 deletions.
diff --git a/src/type_map.rs b/src/type_map.rs
@@ -14,6 +14,15 @@ pub struct Entry<'a, K: 'a, V: 'a> {
 }
 
 impl<'a, K, V> Entry<'a, K, V> {
+    #[inline]
+    pub fn or_insert(self, default: V) -> &'a mut V
+    where
+        V: Send + Sync + 'static,
+    {
+        let v = self.inner.or_insert(Box::new(default));
+        v.downcast_mut().unwrap()
+    }
+
     #[inline]
     pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V
     where
@@ -22,6 +31,36 @@ impl<'a, K, V> Entry<'a, K, V> {
         let v = self.inner.or_insert_with(|| Box::new(default()));
         v.downcast_mut().unwrap()
     }
+
+    #[inline]
+    pub fn or_insert_with_key<F: FnOnce(&K) -> V>(self, default: F) -> &'a mut V
+    where
+        V: Send + Sync + 'static,
+    {
+        let v = self.inner.or_insert_with_key(|key| Box::new(default(key)));
+        v.downcast_mut().unwrap()
+    }
+
+    #[inline]
+    pub fn and_modify<F: FnOnce(&mut V)>(self, f: F) -> Self
+    where
+        V: Send + Sync + 'static,
+    {
+        Entry {
+            inner: self.inner.and_modify(|v| {
+                f(v.downcast_mut().unwrap());
+            }),
+            _marker: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn or_default(self) -> &'a mut V
+    where
+        V: Default + Send + Sync + 'static,
+    {
+        self.or_insert(V::default())
+    }
 }
 
 #[derive(Debug, Default)]
@@ -42,6 +81,13 @@ impl TypeMap {
             .and_then(|boxed| boxed.downcast_ref())
     }
 
+    #[inline]
+    pub fn get_mut<T: 'static>(&mut self) -> Option<&mut T> {
+        self.inner
+            .get_mut(&TypeId::of::<T>())
+            .and_then(|boxed| boxed.downcast_mut())
+    }
+
     #[inline]
     pub fn contains<T: 'static>(&self) -> bool {
         self.inner.contains_key(&TypeId::of::<T>())