1 // Copyright (C) 2024 Huawei Device Co., Ltd. 2 // Licensed under the Apache License, Version 2.0 (the "License"); 3 // you may not use this file except in compliance with the License. 4 // You may obtain a copy of the License at 5 // 6 // http://www.apache.org/licenses/LICENSE-2.0 7 // 8 // Unless required by applicable law or agreed to in writing, software 9 // distributed under the License is distributed on an "AS IS" BASIS, 10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11 // See the License for the specific language governing permissions and 12 // limitations under the License. 13 14 use std::collections::hash_map::Entry; 15 use std::collections::HashMap; 16 use std::hash::Hash; 17 use std::ptr; 18 19 struct Node<K, V> { 20 key: K, 21 value: V, 22 prev: *mut Node<K, V>, 23 next: *mut Node<K, V>, 24 } 25 26 struct LinkedList<K, V> { 27 head: *mut Node<K, V>, 28 tail: *mut Node<K, V>, 29 } 30 31 impl<K, V> LinkedList<K, V> { new() -> Self32 fn new() -> Self { 33 LinkedList { 34 head: ptr::null_mut(), 35 tail: ptr::null_mut(), 36 } 37 } 38 push_front(&mut self, node: *mut Node<K, V>)39 fn push_front(&mut self, node: *mut Node<K, V>) { 40 unsafe { 41 (*node).prev = ptr::null_mut(); 42 (*node).next = self.head; 43 44 if !self.head.is_null() { 45 (*self.head).prev = node; 46 } 47 self.head = node; 48 49 if self.tail.is_null() { 50 self.tail = node; 51 } 52 } 53 } 54 remove(&mut self, node: *mut Node<K, V>)55 fn remove(&mut self, node: *mut Node<K, V>) { 56 unsafe { 57 if !(*node).prev.is_null() { 58 (*(*node).prev).next = (*node).next; 59 } else { 60 self.head = (*node).next; 61 } 62 63 if !(*node).next.is_null() { 64 (*(*node).next).prev = (*node).prev; 65 } else { 66 self.tail = (*node).prev; 67 } 68 } 69 } 70 pop_back(&mut self) -> *mut Node<K, V>71 fn pop_back(&mut self) -> *mut Node<K, V> { 72 if self.tail.is_null() { 73 return ptr::null_mut(); 74 } 75 let node = self.tail; 76 self.remove(node); 77 node 78 } 79 } 80 81 impl<K, V> Drop for LinkedList<K, V> { drop(&mut self)82 fn drop(&mut self) { 83 let mut current = self.head; 84 while !current.is_null() { 85 unsafe { 86 let next = (*current).next; 87 let _ = Box::from_raw(current); 88 current = next; 89 } 90 } 91 } 92 } 93 94 pub struct LRUCache<K, V> { 95 map: HashMap<K, *mut Node<K, V>>, 96 list: LinkedList<K, V>, 97 } 98 99 impl<K: Hash + Eq + Clone, V> LRUCache<K, V> { new() -> Self100 pub fn new() -> Self { 101 LRUCache { 102 map: HashMap::new(), 103 list: LinkedList::new(), 104 } 105 } 106 get(&mut self, key: &K) -> Option<&V>107 pub fn get(&mut self, key: &K) -> Option<&V> { 108 if let Some(&node) = self.map.get(key) { 109 self.list.remove(node); 110 self.list.push_front(node); 111 unsafe { 112 return Some(&(*node).value); 113 } 114 } 115 None 116 } 117 get_mut(&mut self, key: &K) -> Option<&mut V>118 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { 119 if let Some(&mut node) = self.map.get_mut(key) { 120 self.list.remove(node); 121 self.list.push_front(node); 122 unsafe { 123 return Some(&mut (*node).value); 124 } 125 } 126 None 127 } 128 insert(&mut self, key: K, value: V) -> Option<V>129 pub fn insert(&mut self, key: K, value: V) -> Option<V> { 130 match self.map.entry(key) { 131 Entry::Occupied(addr) => { 132 self.list.remove(*addr.get()); 133 self.list.push_front(*addr.get()); 134 unsafe { 135 let old = std::mem::replace(&mut (*(*addr.get())).value, value); 136 Some(old) 137 } 138 } 139 Entry::Vacant(addr) => { 140 let new_node = Box::into_raw(Box::new(Node { 141 key: addr.key().clone(), 142 value, 143 prev: ptr::null_mut(), 144 next: ptr::null_mut(), 145 })); 146 self.list.push_front(new_node); 147 addr.insert(new_node); 148 None 149 } 150 } 151 } 152 pop(&mut self) -> Option<V>153 pub fn pop(&mut self) -> Option<V> { 154 let old_node = self.list.pop_back(); 155 if !old_node.is_null() { 156 unsafe { 157 let old_key = (*old_node).key.clone(); 158 self.map.remove(&old_key); 159 let node = Box::from_raw(old_node); 160 Some(node.value) 161 } 162 } else { 163 None 164 } 165 } 166 remove(&mut self, key: &K) -> Option<V>167 pub fn remove(&mut self, key: &K) -> Option<V> { 168 if let Some(node) = self.map.remove(key) { 169 self.list.remove(node); 170 unsafe { 171 let node = Box::from_raw(node); 172 return Some(node.value); 173 } 174 } 175 None 176 } 177 contains_key(&self, k: &K) -> bool178 pub fn contains_key(&self, k: &K) -> bool { 179 self.map.contains_key(k) 180 } 181 is_empty(&self) -> bool182 pub fn is_empty(&self) -> bool { 183 self.map.is_empty() 184 } 185 len(&self) -> usize186 pub fn len(&self) -> usize { 187 self.map.len() 188 } 189 } 190 191 impl<K: Eq + Hash + Clone, V> Default for LRUCache<K, V> { default() -> Self192 fn default() -> Self { 193 Self::new() 194 } 195 } 196 197 unsafe impl<K, V> Send for LRUCache<K, V> {} 198 199 #[cfg(test)] 200 mod ut_test { 201 use super::LRUCache; 202 203 #[derive(Debug, Eq, PartialEq)] 204 struct Cache { 205 data_count: usize, 206 } 207 208 impl Cache { from_u(init: usize) -> Self209 pub(crate) fn from_u(init: usize) -> Self { 210 Cache { data_count: init } 211 } 212 add(&mut self, num: usize)213 pub(crate) fn add(&mut self, num: usize) { 214 self.data_count += num 215 } 216 } 217 218 #[test] ut_test_empty()219 fn ut_test_empty() { 220 let mut cache = LRUCache::<&str, Cache>::new(); 221 assert_eq!(None, cache.get(&"key1")); 222 assert_eq!(None, cache.get_mut(&"key1")); 223 assert_eq!(0, cache.len()); 224 assert!(!cache.contains_key(&"key1")); 225 assert!(cache.is_empty()); 226 assert_eq!(None, cache.pop()); 227 assert_eq!(None, cache.remove(&"key1")); 228 } 229 230 #[test] ut_test_insert()231 fn ut_test_insert() { 232 let mut cache = LRUCache::new(); 233 assert_eq!(None, cache.insert("key1", Cache::from_u(0))); 234 assert_eq!(Some(&Cache::from_u(0)), cache.get(&"key1")); 235 assert_eq!(Some(&mut Cache::from_u(0)), cache.get_mut(&"key1")); 236 assert_eq!( 237 Some(&mut Cache::from_u(1)), 238 cache.get_mut(&"key1").map(|cache| { 239 cache.add(1); 240 cache 241 }) 242 ); 243 assert_eq!(1, cache.len()); 244 assert!(cache.contains_key(&"key1")); 245 assert!(!cache.is_empty()); 246 assert_eq!(Some(Cache::from_u(1)), cache.pop()); 247 assert_eq!(None, cache.remove(&"key1")); 248 } 249 250 #[test] ut_test_insert_dump()251 fn ut_test_insert_dump() { 252 let mut cache = LRUCache::new(); 253 cache.insert("key0", Cache::from_u(0)); 254 cache.insert("key1", Cache::from_u(1)); 255 cache.insert("key2", Cache::from_u(2)); 256 cache.insert("key3", Cache::from_u(3)); 257 assert_eq!( 258 Some(Cache::from_u(0)), 259 cache.insert("key0", Cache::from_u(4)) 260 ); 261 assert_eq!(4, cache.len()); 262 assert_eq!(Some(Cache::from_u(1)), cache.pop()); 263 assert_eq!(None, cache.get(&"key1")); 264 assert_eq!(Some(&Cache::from_u(4)), cache.get(&"key0")); 265 } 266 267 #[test] ut_test_pop()268 fn ut_test_pop() { 269 let mut cache = LRUCache::new(); 270 cache.insert("key1", Cache::from_u(1)); 271 assert_eq!(Some(Cache::from_u(1)), cache.pop()); 272 assert_eq!(None, cache.get(&"key1")); 273 assert_eq!(None, cache.get_mut(&"key1")); 274 assert_eq!(0, cache.len()); 275 assert!(!cache.contains_key(&"key1")); 276 assert!(cache.is_empty()); 277 assert_eq!(None, cache.pop()); 278 assert_eq!(None, cache.remove(&"key1")); 279 } 280 281 #[test] ut_test_pop_remaining()282 fn ut_test_pop_remaining() { 283 let mut cache = LRUCache::new(); 284 cache.insert("key0", Cache::from_u(0)); 285 cache.insert("key1", Cache::from_u(1)); 286 cache.insert("key2", Cache::from_u(2)); 287 cache.insert("key3", Cache::from_u(3)); 288 assert_eq!(Some(Cache::from_u(0)), cache.pop()); 289 assert_eq!(None, cache.get(&"key0")); 290 assert_eq!(None, cache.get_mut(&"key0")); 291 assert_eq!(3, cache.len()); 292 assert!(!cache.contains_key(&"key0")); 293 assert!(!cache.is_empty()); 294 } 295 296 #[test] ut_test_remove()297 fn ut_test_remove() { 298 let mut cache = LRUCache::new(); 299 cache.insert("key1", Cache::from_u(1)); 300 assert_eq!(Some(Cache::from_u(1)), cache.remove(&"key1")); 301 assert_eq!(None, cache.get(&"key1")); 302 assert_eq!(None, cache.get_mut(&"key1")); 303 assert_eq!(0, cache.len()); 304 assert!(!cache.contains_key(&"key1")); 305 assert!(cache.is_empty()); 306 assert_eq!(None, cache.pop()); 307 assert_eq!(None, cache.remove(&"key1")); 308 } 309 310 #[test] ut_test_remove_remaining()311 fn ut_test_remove_remaining() { 312 let mut cache = LRUCache::new(); 313 cache.insert("key0", Cache::from_u(0)); 314 cache.insert("key1", Cache::from_u(1)); 315 cache.insert("key2", Cache::from_u(2)); 316 cache.insert("key3", Cache::from_u(3)); 317 assert_eq!(Some(Cache::from_u(1)), cache.remove(&"key1")); 318 assert_eq!(None, cache.get(&"key1")); 319 assert_eq!(None, cache.get_mut(&"key1")); 320 assert_eq!(3, cache.len()); 321 assert!(!cache.contains_key(&"key1")); 322 assert!(!cache.is_empty()); 323 assert_eq!(None, cache.remove(&"key1")); 324 } 325 326 #[test] ut_test_insert_after_pop()327 fn ut_test_insert_after_pop() { 328 let mut cache = LRUCache::new(); 329 cache.insert("key0", Cache::from_u(0)); 330 cache.insert("key1", Cache::from_u(1)); 331 assert_eq!(Some(Cache::from_u(0)), cache.pop()); 332 cache.insert("key0", Cache::from_u(0)); 333 assert_eq!(Some(&Cache::from_u(0)), cache.get(&"key0")); 334 assert_eq!(Some(&mut Cache::from_u(0)), cache.get_mut(&"key0")); 335 assert_eq!( 336 Some(&mut Cache::from_u(4)), 337 cache.get_mut(&"key0").map(|cache| { 338 cache.add(4); 339 cache 340 }) 341 ); 342 assert_eq!(Some(&Cache::from_u(4)), cache.get(&"key0")); 343 assert_eq!(2, cache.len()); 344 assert!(cache.contains_key(&"key0")); 345 assert!(!cache.is_empty()); 346 assert_eq!(Some(Cache::from_u(1)), cache.pop()); 347 } 348 } 349