248 lines
6.1 KiB
Rust
248 lines
6.1 KiB
Rust
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<K: Hash + Eq, V: ?Sized, S = RandomState> {
|
|
inner: DashMap<K, Arc<V>, S>,
|
|
}
|
|
|
|
/// A handle that can be converted to a &T or an Arc<T>
|
|
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<T: ?Sized> Deref for ArcRef<'_, T> {
|
|
type Target = Arc<T>;
|
|
fn deref(&self) -> &Self::Target {
|
|
unsafe { std::mem::transmute(self) }
|
|
}
|
|
}
|
|
|
|
impl<'a, T: ?Sized> ArcRef<'a, T> {
|
|
/// Converts the ArcRef into an Arc<T>
|
|
pub fn to_arc(self) -> Arc<T> {
|
|
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<K: Hash + Eq, V: ?Sized, S: BuildHasher + Default + Clone> Default for CacheMap<K, V, S> {
|
|
fn default() -> Self {
|
|
CacheMap {
|
|
inner: Default::default(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<K: Hash + Eq, V, S: BuildHasher + Default + Clone> std::iter::FromIterator<(K, V)>
|
|
for CacheMap<K, V, S>
|
|
{
|
|
fn from_iter<T>(iter: T) -> Self
|
|
where
|
|
T: IntoIterator<Item = (K, V)>,
|
|
{
|
|
CacheMap {
|
|
inner: iter.into_iter().map(|(k, v)| (k, Arc::new(v))).collect(),
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct IntoIter<K, V, S>(dashmap::iter::OwningIter<K, Arc<V>, S>);
|
|
|
|
impl<K: Eq + Hash, V, S: BuildHasher + Clone> Iterator for IntoIter<K, V, S> {
|
|
type Item = (K, Arc<V>);
|
|
|
|
fn next(&mut self) -> Option<Self::Item> {
|
|
self.0.next()
|
|
}
|
|
}
|
|
|
|
impl<K: Hash + Eq, V, S: BuildHasher + Clone> IntoIterator for CacheMap<K, V, S> {
|
|
type Item = (K, Arc<V>);
|
|
type IntoIter = IntoIter<K, V, S>;
|
|
|
|
fn into_iter(self) -> Self::IntoIter {
|
|
IntoIter(self.inner.into_iter())
|
|
}
|
|
}
|
|
|
|
impl<K: Hash + Eq, V, S: BuildHasher + Clone> CacheMap<K, V, S> {
|
|
/// 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<F: FnOnce() -> 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<Q: ?Sized>(&self, key: &Q) -> bool
|
|
where
|
|
K: std::borrow::Borrow<Q>,
|
|
Q: Hash + Eq,
|
|
{
|
|
self.inner.contains_key(key)
|
|
}
|
|
}
|
|
|
|
impl<K: Hash + Eq, V: ?Sized> CacheMap<K, V, RandomState> {
|
|
/// Creates a new CacheMap
|
|
pub fn new() -> Self {
|
|
CacheMap {
|
|
inner: DashMap::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<K: Hash + Eq, V: ?Sized, S: BuildHasher + Clone> CacheMap<K, V, S> {
|
|
/// 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<V>, 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<F: FnOnce() -> Arc<V>>(&self, key: K, f: F) -> ArcRef<'_, V> {
|
|
let val = self.inner.entry(key).or_insert_with(f);
|
|
let arc: &Arc<V> = &*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);
|
|
}
|
|
}
|