reifydb_core/util/
lru.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use std::{
5	hash::Hash,
6	mem,
7	sync::atomic::{AtomicU64, Ordering},
8};
9
10use dashmap::DashMap;
11
12pub struct LruCache<K, V> {
13	map: DashMap<K, Entry<V>>,
14	capacity: usize,
15	access_counter: AtomicU64,
16}
17
18struct Entry<V> {
19	value: V,
20	last_access: u64,
21}
22
23impl<K: Hash + Eq + Clone, V: Clone> LruCache<K, V> {
24	pub fn new(capacity: usize) -> Self {
25		assert!(capacity > 0, "LRU cache capacity must be greater than 0");
26		Self {
27			map: DashMap::with_capacity(capacity),
28			capacity,
29			access_counter: AtomicU64::new(0),
30		}
31	}
32
33	pub fn get(&self, key: &K) -> Option<V> {
34		if let Some(mut entry) = self.map.get_mut(key) {
35			entry.last_access = self.access_counter.fetch_add(1, Ordering::Relaxed);
36			Some(entry.value.clone())
37		} else {
38			None
39		}
40	}
41
42	pub fn put(&self, key: K, value: V) -> Option<V> {
43		let access = self.access_counter.fetch_add(1, Ordering::Relaxed);
44
45		if let Some(mut entry) = self.map.get_mut(&key) {
46			let old_value = mem::replace(&mut entry.value, value);
47			entry.last_access = access;
48			return Some(old_value);
49		}
50
51		if self.map.len() >= self.capacity {
52			self.evict_lru();
53		}
54
55		self.map.insert(
56			key,
57			Entry {
58				value,
59				last_access: access,
60			},
61		);
62		None
63	}
64
65	pub fn remove(&self, key: &K) -> Option<V> {
66		self.map.remove(key).map(|(_, entry)| entry.value)
67	}
68
69	pub fn contains_key(&self, key: &K) -> bool {
70		self.map.contains_key(key)
71	}
72
73	pub fn clear(&self) {
74		self.map.clear();
75	}
76
77	pub fn len(&self) -> usize {
78		self.map.len()
79	}
80
81	pub fn is_empty(&self) -> bool {
82		self.map.is_empty()
83	}
84
85	pub fn capacity(&self) -> usize {
86		self.capacity
87	}
88
89	fn evict_lru(&self) {
90		let mut oldest_key: Option<K> = None;
91		let mut oldest_access = u64::MAX;
92
93		for entry in self.map.iter() {
94			if entry.last_access < oldest_access {
95				oldest_access = entry.last_access;
96				oldest_key = Some(entry.key().clone());
97			}
98		}
99
100		if let Some(key) = oldest_key {
101			self.map.remove(&key);
102		}
103	}
104}
105
106#[cfg(test)]
107pub mod tests {
108
109	mod lru {
110		use crate::util::lru::LruCache;
111
112		#[test]
113		fn test_basic_operations() {
114			let cache = LruCache::new(2);
115
116			assert_eq!(cache.put(1, "a"), None);
117			assert_eq!(cache.put(2, "b"), None);
118			assert_eq!(cache.get(&1), Some("a"));
119			assert_eq!(cache.get(&2), Some("b"));
120			assert_eq!(cache.len(), 2);
121		}
122
123		#[test]
124		fn test_eviction() {
125			let cache = LruCache::new(2);
126
127			cache.put(1, "a");
128			cache.put(2, "b");
129			let evicted = cache.put(3, "c");
130
131			// Eviction returns None because it's a new key insertion
132			// The evicted value is the LRU entry (key 1)
133			assert_eq!(evicted, None);
134			assert_eq!(cache.get(&1), None);
135			assert_eq!(cache.get(&2), Some("b"));
136			assert_eq!(cache.get(&3), Some("c"));
137		}
138
139		#[test]
140		fn test_lru_order() {
141			let cache = LruCache::new(2);
142
143			cache.put(1, "a");
144			cache.put(2, "b");
145			cache.get(&1); // Access 1, making it more recent than 2
146			cache.put(3, "c"); // Should evict 2 (least recently used)
147
148			assert_eq!(cache.get(&1), Some("a"));
149			assert_eq!(cache.get(&2), None);
150			assert_eq!(cache.get(&3), Some("c"));
151		}
152
153		#[test]
154		fn test_update_existing() {
155			let cache = LruCache::new(2);
156
157			cache.put(1, "a");
158			let old = cache.put(1, "b");
159
160			assert_eq!(old, Some("a"));
161			assert_eq!(cache.get(&1), Some("b"));
162			assert_eq!(cache.len(), 1);
163		}
164
165		#[test]
166		fn test_remove() {
167			let cache = LruCache::new(2);
168
169			cache.put(1, "a");
170			cache.put(2, "b");
171			let removed = cache.remove(&1);
172
173			assert_eq!(removed, Some("a"));
174			assert_eq!(cache.get(&1), None);
175			assert_eq!(cache.len(), 1);
176		}
177
178		#[test]
179		fn test_clear() {
180			let cache = LruCache::new(2);
181
182			cache.put(1, "a");
183			cache.put(2, "b");
184			cache.clear();
185
186			assert_eq!(cache.len(), 0);
187			assert!(cache.is_empty());
188		}
189
190		#[test]
191		fn test_contains_key() {
192			let cache = LruCache::new(2);
193
194			cache.put(1, "a");
195			assert!(cache.contains_key(&1));
196			assert!(!cache.contains_key(&2));
197		}
198	}
199}
reifydb_core/util/lru.rs

reifydb_core/util/
lru.rs
