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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

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

use core::cmp::Eq;
use core::hash::BuildHasher;
use core::hash::Hash;
use core::ops::Index;
use std::sync::Arc;

use crate::{FnCache, FnCacheMany};

/// 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<'f, I, O, S = RandomState>
where
	I: Eq + Hash,
{
	pub(crate) cache: HashMap<I, O, S>,
	f: Arc<dyn Fn(&mut Self, &I) -> O + 'f + Send + Sync>,
}

impl<'f, I, O, S> FnCache<I, O> for HashCache<'f, 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)
		}
	}
}

impl<'f, I, O, S> FnCacheMany<I, O> for HashCache<'f, I, O, S>
where
	I: Eq + Hash + Clone,
	S: BuildHasher,
{
	fn get_many<const N: usize>(&mut self, inputs: [I; N]) -> [&O; N] {
		for i in &inputs {
			self.get(i.clone());
		}

		inputs.map(|i| self.cache.get(&i).unwrap())
	}
}

impl<'f, I, O> HashCache<'f, 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(&I) -> O + 'f + Send + Sync,
	{
		Self::recursive(move |_, i| f(i))
	}

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

impl<'f, I, O, S> HashCache<'f, 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::with_hasher`] for more details.
	pub fn with_hasher<F>(hash_builder: S, f: F) -> Self
	where
		F: Fn(&I) -> O + 'f + Send + Sync,
	{
		Self::recursive_with_hasher(hash_builder, move |_, i| f(i))
	}

	/// Create a recursive HashCache which will use the given hash
	/// builder to hash keys.
	///
	/// See the documentation on [`HashMap::with_hasher`] for more details.
	pub fn recursive_with_hasher<F>(hash_builder: S, f: F) -> Self
	where
		F: Fn(&mut Self, &I) -> O + 'f + Send + Sync,
	{
		HashCache {
			cache: HashMap::with_hasher(hash_builder),
			f: Arc::new(f),
		}
	}

	fn compute(&mut self, input: &I) -> O {
		(self.f.clone())(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)
	}
}