use dashmap::DashMap; use std::ops::Deref; use std::sync::Arc; use std::{collections::hash_map::RandomState, hash::Hash}; use std::{hash::BuildHasher, marker::PhantomData}; /// An insert-only map for caching the result of functions pub struct CacheMap { inner: DashMap, S>, } /// A handle that can be converted to a &T or an Arc pub struct ArcRef<'a, T: ?Sized> { // this pointer never gets dereferenced, but it has to be T, so that Ref is the right size for wide pointers #[allow(dead_code)] fake_ptr: *const T, phantom: PhantomData<&'a T>, } impl<'a, T: ?Sized> Clone for ArcRef<'a, T> { fn clone(&self) -> Self { *self } } impl<'a, T: ?Sized> Copy for ArcRef<'a, T> {} impl Deref for ArcRef<'_, T> { type Target = Arc; fn deref(&self) -> &Self::Target { unsafe { std::mem::transmute(self) } } } impl<'a, T: ?Sized> ArcRef<'a, T> { /// Converts the ArcRef into an Arc pub fn to_arc(self) -> Arc { self.deref().clone() } /// Converts the ArcRef into a &T pub fn as_ref(self) -> &'a T { let ptr = &**self as *const T; unsafe { &*ptr } } } impl Default for CacheMap { fn default() -> Self { CacheMap { inner: Default::default(), } } } impl std::iter::FromIterator<(K, V)> for CacheMap { fn from_iter(iter: T) -> Self where T: IntoIterator, { CacheMap { inner: iter.into_iter().map(|(k, v)| (k, Arc::new(v))).collect(), } } } pub struct IntoIter(dashmap::iter::OwningIter, S>); impl Iterator for IntoIter { type Item = (K, Arc); fn next(&mut self) -> Option { self.0.next() } } impl IntoIterator for CacheMap { type Item = (K, Arc); type IntoIter = IntoIter; fn into_iter(self) -> Self::IntoIter { IntoIter(self.inner.into_iter()) } } impl CacheMap { /// Fetch the value associated with the key, or run the provided function to insert one. /// /// # Example /// /// ``` /// use cachemap2::CacheMap; /// /// let m = CacheMap::new(); /// /// let fst = m.cache("key", || 5u32).as_ref(); /// let snd = m.cache("key", || 7u32).as_ref(); /// /// assert_eq!(*fst, *snd); /// assert_eq!(*fst, 5u32); /// ``` pub fn cache V>(&self, key: K, f: F) -> ArcRef<'_, V> { self.cache_arc(key, || Arc::new(f())) } /// Fetch the value associated with the key, or insert a default value. pub fn cache_default(&self, key: K) -> ArcRef<'_, V> where V: Default, { self.cache(key, || Default::default()) } /// Return whether the map contains the given key. pub fn contains_key(&self, key: &Q) -> bool where K: std::borrow::Borrow, Q: Hash + Eq, { self.inner.contains_key(key) } } impl CacheMap { /// Creates a new CacheMap pub fn new() -> Self { CacheMap { inner: DashMap::new(), } } } impl CacheMap { /// Creates a new CacheMap with the provided hasher pub fn with_hasher(hash_builder: S) -> Self { Self { inner: DashMap::with_hasher(hash_builder), } } /// Fetch the value associated with the key, or run the provided function to insert one. /// With this version, the function returns an Arc, whch allows caching unsized types. /// /// # Example /// /// ``` /// use cachemap2::CacheMap; /// /// let m: CacheMap<_, [usize]> = CacheMap::new(); /// /// let a = m.cache_arc("a", || { /// let a = &[1,2,3][..]; /// a.into() /// }).as_ref(); /// /// let b = m.cache_arc("b", || { /// let b = &[9,9][..]; /// b.into() /// }).as_ref(); /// /// assert_eq!(a, &[1,2,3]); /// assert_eq!(b, &[9,9]); /// ``` pub fn cache_arc Arc>(&self, key: K, f: F) -> ArcRef<'_, V> { let val = self.inner.entry(key).or_insert_with(f); let arc: &Arc = &*val; let arc_ref: &ArcRef<'_, V> = unsafe { std::mem::transmute(arc) }; *arc_ref } } #[cfg(test)] mod tests { use super::*; #[test] fn single_insert() { let m = CacheMap::new(); let a = m.cache("key", || 21u32).as_ref(); assert_eq!(21, *a); } #[test] fn contains_key() { let m = CacheMap::new(); m.cache("key", || 21u32); assert!(m.contains_key("key")); assert!(!m.contains_key("other")); } #[test] fn double_insert() { let m = CacheMap::new(); let a = m.cache("key", || 5u32).as_ref(); let b = m.cache("key", || 7u32).as_ref(); assert_eq!(*a, *b); assert_eq!(5, *a); } #[test] fn insert_two() { let m = CacheMap::new(); let a = m.cache("a", || 5u32).as_ref(); let b = m.cache("b", || 7u32).as_ref(); assert_eq!(5, *a); assert_eq!(7, *b); let c = m.cache("a", || 9u32).as_ref(); let d = m.cache("b", || 11u32).as_ref(); assert_eq!(*a, *c); assert_eq!(*b, *d); assert_eq!(5, *a); assert_eq!(7, *b); } #[test] fn use_after_drop() { #[derive(Clone)] struct Foo(usize); impl Drop for Foo { fn drop(&mut self) { assert_eq!(33, self.0); } } { let mut arc = { let m = CacheMap::new(); let a = m.cache("key", || Foo(99)).to_arc(); assert_eq!(99, (*a).0); a }; Arc::make_mut(&mut arc).0 = 33; } assert!(true); } }