1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

use std::collections::HashMap;
use std::collections::hash_map::RandomState;

use core::cmp::Eq;
use core::hash::Hash;
use core::hash::BuildHasher;
use core::ops::Index;

use crate::FnCache;

/// A cache for a function which uses a `HashMap`.
///
/// The cache takes ownership of all inputs, but
/// only passes a reference to the function,
/// allowing it to store the input in the cache
/// without any copies or clones.
///
/// The requirements for a `HashMap` must be met,
/// specifically the keys must implement `Eq` and
/// `Hash`, and the following propery must hold:
///
/// ```k1 == k2 -> hash(k1) == hash(k2)```
pub struct HashCache<'a, I, O, S = RandomState>
where
	I: Eq + Hash,
{
	pub(crate) cache: HashMap<I, O, S>,
	f: *mut (dyn Fn(&mut Self, &I) -> O + 'a),
}

impl<'a, I, O, S> FnCache<I, O> for HashCache<'a, I, O, S>
where
	I: Eq + Hash,
	S: BuildHasher,
{
	fn get(&mut self, input: I) -> &O {
		if self.cache.contains_key(&input) {
			self.cache.index(&input)
		} else {
			let output = self.compute(&input);
			self.cache.entry(input).or_insert(output.into())
		}
	}
}

impl<'a, I, O> HashCache<'a, I, O, RandomState>
where
	I: Eq + Hash,
{
	/// Create a cache for the provided function. If the
	/// function stores references, the cache can only
	/// live as long as those references.
	pub fn new<F>(f: F) -> Self
	where
		F: Fn(&mut Self, &I) -> O + 'a,
	{
		HashCache {
			cache: HashMap::default(),
			f: Box::into_raw(Box::new(f)),
		}
	}
}


impl<'a, I, O, S> HashCache<'a, I, O, S>
where
	I: Eq + Hash,
	S: BuildHasher,
{
	/// Create a HashCache which will use the given hash
	/// builder to hash keys.
	///
	/// See the documentation on `HashMap` for more details.
	pub fn with_hasher<F>(hash_builder: S, f: F) -> Self
	where
		F: Fn(&mut Self, &I) -> O + 'a,
	{
		HashCache {
			cache: HashMap::with_hasher(hash_builder),
			f: Box::into_raw(Box::new(f)),
		}
	}

	fn compute(&mut self, input: &I) -> O {
		unsafe { (*self.f)(self, input) }
	}

	/// Clears the cache, removing all key-value pairs.
	/// Keeps the allocated memory for reuse.
	pub fn clear(&mut self) {
		self.cache.clear();
	}

	/// Returns the number of elements in the cache.
	pub fn len(&self) -> usize {
		self.cache.len()
	}

	/// Reserves capacity for at least `additional` more elements
	/// to be inserted in the cache. The collection may
	/// reserve more space to avoid frequent reallocations.
	pub fn reserve(&mut self, additional: usize) {
		self.cache.reserve(additional)
	}

	/// Removes the input from the cache, returning any value
	/// if the input was previously in the cache.
	pub fn remove(&mut self, input: &I) -> Option<O> {
		self.cache.remove(input)
	}
}

#[doc(hidden)]
impl<'a, I, O, S> Drop for HashCache<'a, I, O, S>
where
	I: Eq + Hash,
{
	fn drop(&mut self) {
		#[allow(unused_must_use)]
		unsafe {
			Box::from_raw(self.f);
		}
	}
}