1use std::collections::hash_map::Entry;
2use std::collections::{HashMap, VecDeque};
3use std::hash::Hash;
4
5pub struct TwoQCache<K, V> {
6 hot: HashMap<K, V>,
7 hot_order: VecDeque<K>,
8 cold: HashMap<K, V>,
9 cold_order: VecDeque<K>,
10 hot_limit: usize,
11 cold_limit: usize,
12}
13
14impl<K: Hash + Eq + Clone, V> TwoQCache<K, V> {
15 pub fn new(hot_limit: usize, cold_limit: usize) -> Self {
16 Self {
17 hot: HashMap::new(),
18 hot_order: VecDeque::new(),
19 cold: HashMap::new(),
20 cold_order: VecDeque::new(),
21 hot_limit,
22 cold_limit,
23 }
24 }
25
26 pub fn get(&mut self, key: &K) -> Option<&V> {
27 if self.hot.contains_key(key) {
28 if let Some(pos) = self.hot_order.iter().position(|k| k == key) {
29 let k = self.hot_order.remove(pos).unwrap();
30 self.hot_order.push_back(k);
31 }
32 return self.hot.get(key);
33 }
34
35 if let Some(v) = self.cold.remove(key) {
36 if let Some(pos) = self.cold_order.iter().position(|k| k == key) {
37 self.cold_order.remove(pos);
38 }
39 self.hot.insert(key.clone(), v);
40 self.hot_order.push_back(key.clone());
41 self.evict_hot();
42 return self.hot.get(key);
43 }
44
45 None
46 }
47
48 pub fn insert(&mut self, key: K, value: V) {
49 match self.hot.entry(key) {
50 Entry::Occupied(mut e) => {
51 e.insert(value);
52 }
53 Entry::Vacant(e) => {
54 let key = e.into_key();
55 match self.cold.entry(key) {
56 Entry::Occupied(mut e) => {
57 e.insert(value);
58 }
59 Entry::Vacant(e) => {
60 let key = e.into_key();
61 self.cold.insert(key.clone(), value);
62 self.cold_order.push_back(key);
63 self.evict_cold();
64 }
65 }
66 }
67 }
68 }
69
70 #[inline]
71 #[must_use]
72 pub fn contains(&self, key: &K) -> bool {
73 self.hot.contains_key(key) || self.cold.contains_key(key)
74 }
75
76 #[inline]
77 #[must_use]
78 pub fn len(&self) -> usize {
79 self.hot.len() + self.cold.len()
80 }
81
82 #[inline]
83 #[must_use]
84 pub fn is_empty(&self) -> bool {
85 self.hot.is_empty() && self.cold.is_empty()
86 }
87
88 #[inline]
89 #[must_use]
90 pub fn hot_len(&self) -> usize {
91 self.hot.len()
92 }
93
94 #[inline]
95 #[must_use]
96 pub fn cold_len(&self) -> usize {
97 self.cold.len()
98 }
99
100 pub fn remove(&mut self, key: &K) -> Option<V> {
101 if let Some(v) = self.hot.remove(key) {
102 if let Some(pos) = self.hot_order.iter().position(|k| k == key) {
103 self.hot_order.remove(pos);
104 }
105 return Some(v);
106 }
107 if let Some(v) = self.cold.remove(key) {
108 if let Some(pos) = self.cold_order.iter().position(|k| k == key) {
109 self.cold_order.remove(pos);
110 }
111 return Some(v);
112 }
113 None
114 }
115
116 fn evict_hot(&mut self) {
117 while self.hot.len() > self.hot_limit {
118 if let Some(oldest) = self.hot_order.pop_front()
119 && let Some(v) = self.hot.remove(&oldest)
120 {
121 self.cold.insert(oldest.clone(), v);
122 self.cold_order.push_back(oldest);
123 self.evict_cold();
124 }
125 }
126 }
127
128 fn evict_cold(&mut self) {
129 while self.cold.len() > self.cold_limit {
130 if let Some(oldest) = self.cold_order.pop_front() {
131 self.cold.remove(&oldest);
132 }
133 }
134 }
135}
136
137#[cfg(test)]
138mod tests {
139 use super::*;
140
141 #[test]
142 fn insert_and_get() {
143 let mut cache = TwoQCache::new(2, 4);
144 cache.insert("a", 1);
145 cache.insert("b", 2);
146 assert_eq!(cache.get(&"a"), Some(&1));
147 assert_eq!(cache.get(&"b"), Some(&2));
148 }
149
150 #[test]
151 fn cold_to_hot_promotion() {
152 let mut cache = TwoQCache::new(1, 3);
153 cache.insert("a", 1);
154 assert_eq!(cache.hot_len(), 0);
155 assert_eq!(cache.cold_len(), 1);
156
157 let _ = cache.get(&"a");
158 assert_eq!(cache.hot_len(), 1);
159 assert_eq!(cache.cold_len(), 0);
160 }
161
162 #[test]
163 fn hot_eviction_to_cold() {
164 let mut cache = TwoQCache::new(1, 2);
165 cache.insert("a", 1);
166 let _ = cache.get(&"a"); cache.insert("b", 2);
168 let _ = cache.get(&"b"); assert_eq!(cache.hot_len(), 1);
171 assert!(cache.contains(&"a")); assert!(cache.contains(&"b")); }
174
175 #[test]
176 fn cold_eviction_drops() {
177 let mut cache = TwoQCache::new(2, 2);
178 cache.insert("a", 1);
179 cache.insert("b", 2);
180 cache.insert("c", 3);
181 cache.insert("d", 4);
182 cache.insert("e", 5);
183
184 assert!(!cache.contains(&"a"));
185 assert!(!cache.contains(&"b"));
186 assert!(!cache.contains(&"c"));
187 assert!(cache.contains(&"d"));
188 assert!(cache.contains(&"e"));
189 assert_eq!(cache.len(), 2);
190 }
191
192 #[test]
193 fn scan_resistance() {
194 let mut cache: TwoQCache<String, i32> = TwoQCache::new(2, 3);
195 cache.insert("hot1".into(), 1);
196 cache.insert("hot2".into(), 2);
197 let _ = cache.get(&"hot1".to_string());
198 let _ = cache.get(&"hot2".to_string());
199
200 for i in 0..10 {
201 cache.insert(format!("scan_{i}"), i);
202 }
203
204 assert!(cache.contains(&"hot1".to_string()));
205 assert!(cache.contains(&"hot2".to_string()));
206 }
207
208 #[test]
209 fn len_and_empty() {
210 let mut cache: TwoQCache<&str, i32> = TwoQCache::new(2, 2);
211 assert!(cache.is_empty());
212 cache.insert("a", 1);
213 assert_eq!(cache.len(), 1);
214 assert!(!cache.is_empty());
215 }
216
217 #[test]
218 fn remove_key() {
219 let mut cache = TwoQCache::new(2, 2);
220 cache.insert("a", 1);
221 assert_eq!(cache.remove(&"a"), Some(1));
222 assert!(!cache.contains(&"a"));
223 assert_eq!(cache.remove(&"a"), None);
224 }
225
226 #[test]
227 fn update_existing_key() {
228 let mut cache = TwoQCache::new(2, 2);
229 cache.insert("a", 1);
230 cache.insert("a", 42);
231 assert_eq!(cache.get(&"a"), Some(&42));
232 }
233
234 #[test]
235 fn missing_get_returns_none() {
236 let mut cache: TwoQCache<&str, i32> = TwoQCache::new(2, 2);
237 assert_eq!(cache.get(&"missing"), None);
238 }
239}