Skip to content

Commit 54d476f

Browse files
obycodeobycode
committed
feat: make LruCache methods fallible
1 parent 2e683a3 commit 54d476f

File tree

1 file changed

+93
-71
lines changed

1 file changed

+93
-71
lines changed

‎stacks-common/src/util/lru_cache.rs

Lines changed: 93 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright (C) 2024 Stacks Open Internet Foundation
1+
// Copyright (C) 2025 Stacks Open Internet Foundation
22
//
33
// This program is free software: you can redistribute it and/or modify
44
// it under the terms of the GNU General Public License as published by
@@ -62,8 +62,12 @@ impl<K: Display, V: Display> Display for LruCache<K, V> {
6262
)?;
6363
let mut curr = self.head;
6464
while curr != self.capacity {
65-
writeln!(f, " {}", self.order[curr])?;
66-
curr = self.order[curr].next;
65+
let Some(node) = self.order.get(curr) else {
66+
writeln!(f, " <invalid>")?;
67+
break;
68+
};
69+
writeln!(f, " {}", node)?;
70+
curr = node.next;
6771
}
6872
Ok(())
6973
}
@@ -82,89 +86,105 @@ impl<K: Eq + std::hash::Hash + Clone, V: Copy> LruCache<K, V> {
8286
}
8387

8488
/// Get the value for the given key
85-
pub fn get(&mut self, key: &K) -> Option<V> {
86-
if let Some(node) = self.cache.get(key) {
89+
/// Returns an error iff the cache is corrupted and should be discarded
90+
pub fn get(&mut self, key: &K) -> Result<Option<V>, ()> {
91+
if let Some(order_idx) = self.cache.get(key) {
8792
// Move the node to the head of the LRU list
88-
let node = *node;
89-
90-
if node != self.head {
91-
let prev = self.order[node].prev;
92-
let next = self.order[node].next;
93-
94-
if node == self.tail {
93+
if *order_idx != self.head {
94+
let node = self.order.get_mut(*order_idx).ok_or(())?;
95+
let prev = node.prev;
96+
let next = node.next;
97+
node.prev = self.capacity;
98+
node.next = self.head;
99+
100+
if *order_idx == self.tail {
95101
// If this is the tail, update the tail
96102
self.tail = prev;
97103
} else {
98104
// Else, update the next node's prev pointer
99-
self.order[next].prev = prev;
105+
let next_node = self.order.get_mut(next).ok_or(())?;
106+
next_node.prev = prev;
100107
}
101108

102-
self.order[prev].next = next;
103-
self.order[node].prev = self.capacity;
104-
self.order[node].next = self.head;
105-
self.order[self.head].prev = node;
106-
self.head = node;
109+
let prev_node = self.order.get_mut(prev).ok_or(())?;
110+
prev_node.next = next;
111+
112+
let head_node = self.order.get_mut(self.head).ok_or(())?;
113+
head_node.prev = *order_idx;
114+
self.head = *order_idx;
107115
}
108116

109-
Some(self.order[node].value)
117+
let node = self.order.get(*order_idx).ok_or(())?;
118+
Ok(Some(node.value))
110119
} else {
111-
None
120+
Ok(None)
112121
}
113122
}
114123

115124
/// Insert a key-value pair into the cache, marking it as dirty.
116-
/// Returns `Some((K, V))` if a dirty value was evicted.
117-
pub fn insert(&mut self, key: K, value: V) -> Option<(K, V)> {
125+
/// Returns an error iff the cache is corrupted and should be discarded
126+
/// Returns `Ok(Some((K, V)))` if a dirty value was evicted.
127+
pub fn insert(&mut self, key: K, value: V) -> Result<Option<(K, V)>, ()> {
118128
self.insert_with_dirty(key, value, true)
119129
}
120130

121131
/// Insert a key-value pair into the cache, marking it as clean.
122-
/// Returns `Some((K, V))` if a dirty value was evicted.
123-
pub fn insert_clean(&mut self, key: K, value: V) -> Option<(K, V)> {
132+
/// Returns an error iff the cache is corrupted and should be discarded
133+
/// Returns `Ok(Some((K, V)))` if a dirty value was evicted.
134+
pub fn insert_clean(&mut self, key: K, value: V) -> Result<Option<(K, V)>, ()> {
124135
self.insert_with_dirty(key, value, false)
125136
}
126137

127138
/// Insert a key-value pair into the cache
128-
/// Returns `Some((K, V))` if a dirty value was evicted.
129-
pub fn insert_with_dirty(&mut self, key: K, value: V, dirty: bool) -> Option<(K, V)> {
139+
/// Returns an error iff the cache is corrupted and should be discarded
140+
/// Returns `Ok(Some((K, V)))` if a dirty value was evicted.
141+
pub fn insert_with_dirty(
142+
&mut self,
143+
key: K,
144+
value: V,
145+
dirty: bool,
146+
) -> Result<Option<(K, V)>, ()> {
130147
let mut evicted = None;
131-
if let Some(node) = self.cache.get(&key) {
148+
if let Some(order_idx) = self.cache.get(&key) {
132149
// Update the value for the key
133-
let node = *node;
134-
self.order[node].value = value;
135-
self.order[node].dirty = dirty;
150+
let node = self.order.get_mut(*order_idx).ok_or(())?;
151+
node.value = value;
152+
node.dirty = dirty;
136153

137154
// Just call get to handle updating the LRU list
138-
self.get(&key);
155+
self.get(&key)?;
139156
} else {
140157
let index = if self.cache.len() == self.capacity {
141158
// Take the place of the least recently used element.
142159
// First, remove it from the tail of the LRU list
143160
let index = self.tail;
144-
let prev = self.order[index].prev;
145-
self.order[prev].next = self.capacity;
146-
self.tail = prev;
161+
let tail_node = self.order.get_mut(index).ok_or(())?;
162+
let prev = tail_node.prev;
147163

148164
// Remove it from the cache
149-
self.cache.remove(&self.order[index].key);
165+
self.cache.remove(&tail_node.key);
150166

151167
// Replace the key with the new key, saving the old key
152-
let replaced_key = std::mem::replace(&mut self.order[index].key, key.clone());
168+
let replaced_key = std::mem::replace(&mut tail_node.key, key.clone());
153169

154170
// If it is dirty, save the key-value pair to return
155-
if self.order[index].dirty {
156-
evicted = Some((replaced_key, self.order[index].value));
171+
if tail_node.dirty {
172+
evicted = Some((replaced_key, tail_node.value));
157173
}
158174

159175
// Insert this new value into the cache
160176
self.cache.insert(key, index);
161177

162178
// Update the node with the new key-value pair, inserting it at
163179
// the head of the LRU list
164-
self.order[index].value = value;
165-
self.order[index].dirty = dirty;
166-
self.order[index].next = self.head;
167-
self.order[index].prev = self.capacity;
180+
tail_node.value = value;
181+
tail_node.dirty = dirty;
182+
tail_node.next = self.head;
183+
tail_node.prev = self.capacity;
184+
185+
let tail_prev_node = self.order.get_mut(prev).ok_or(())?;
186+
tail_prev_node.next = self.capacity;
187+
self.tail = prev;
168188

169189
index
170190
} else {
@@ -193,9 +213,11 @@ impl<K: Eq + std::hash::Hash + Clone, V: Copy> LruCache<K, V> {
193213

194214
self.head = index;
195215
}
196-
evicted
216+
Ok(evicted)
197217
}
198218

219+
/// Flush all dirty values in the cache, calling the given function, `f`,
220+
/// for each dirty value.
199221
pub fn flush<E>(&mut self, mut f: impl FnMut(&K, V) -> Result<(), E>) -> Result<(), E> {
200222
let mut index = self.head;
201223
while index != self.capacity {
@@ -219,47 +241,47 @@ mod tests {
219241
fn test_lru_cache() {
220242
let mut cache = LruCache::new(2);
221243

222-
cache.insert(1, 1);
223-
cache.insert(2, 2);
224-
assert_eq!(cache.get(&1), Some(1));
225-
cache.insert(3, 3);
226-
assert_eq!(cache.get(&2), None);
227-
cache.insert(4, 4);
228-
assert_eq!(cache.get(&1), None);
229-
assert_eq!(cache.get(&3), Some(3));
230-
assert_eq!(cache.get(&4), Some(4));
244+
cache.insert(1, 1).unwrap();
245+
cache.insert(2, 2).unwrap();
246+
assert_eq!(cache.get(&1).unwrap(), Some(1));
247+
cache.insert(3, 3).unwrap();
248+
assert_eq!(cache.get(&2).unwrap(), None);
249+
cache.insert(4, 4).unwrap();
250+
assert_eq!(cache.get(&1).unwrap(), None);
251+
assert_eq!(cache.get(&3).unwrap(), Some(3));
252+
assert_eq!(cache.get(&4).unwrap(), Some(4));
231253
}
232254

233255
#[test]
234256
fn test_lru_cache_update() {
235257
let mut cache = LruCache::new(2);
236258

237-
cache.insert(1, 1);
238-
cache.insert(2, 2);
239-
cache.insert(1, 10);
240-
assert_eq!(cache.get(&1), Some(10));
241-
cache.insert(3, 3);
242-
assert_eq!(cache.get(&2), None);
243-
cache.insert(2, 4);
244-
assert_eq!(cache.get(&2), Some(4));
245-
assert_eq!(cache.get(&3), Some(3));
259+
cache.insert(1, 1).unwrap();
260+
cache.insert(2, 2).unwrap();
261+
cache.insert(1, 10).unwrap();
262+
assert_eq!(cache.get(&1).unwrap(), Some(10));
263+
cache.insert(3, 3).unwrap();
264+
assert_eq!(cache.get(&2).unwrap(), None);
265+
cache.insert(2, 4).unwrap();
266+
assert_eq!(cache.get(&2).unwrap(), Some(4));
267+
assert_eq!(cache.get(&3).unwrap(), Some(3));
246268
}
247269

248270
#[test]
249271
fn test_lru_cache_evicted() {
250272
let mut cache = LruCache::new(2);
251273

252-
assert!(cache.insert(1, 1).is_none());
253-
assert!(cache.insert(2, 2).is_none());
254-
let evicted = cache.insert(3, 3).expect("expected an eviction");
274+
assert!(cache.insert(1, 1).unwrap().is_none());
275+
assert!(cache.insert(2, 2).unwrap().is_none());
276+
let evicted = cache.insert(3, 3).unwrap().expect("expected an eviction");
255277
assert_eq!(evicted, (1, 1));
256278
}
257279

258280
#[test]
259281
fn test_lru_cache_flush() {
260282
let mut cache = LruCache::new(2);
261283

262-
cache.insert(1, 1);
284+
cache.insert(1, 1).unwrap();
263285

264286
let mut flushed = Vec::new();
265287
cache
@@ -271,8 +293,8 @@ mod tests {
271293

272294
assert_eq!(flushed, vec![(1, 1)]);
273295

274-
cache.insert(1, 3);
275-
cache.insert(2, 2);
296+
cache.insert(1, 3).unwrap();
297+
cache.insert(2, 2).unwrap();
276298

277299
let mut flushed = Vec::new();
278300
cache
@@ -289,10 +311,10 @@ mod tests {
289311
fn test_lru_cache_evict_clean() {
290312
let mut cache = LruCache::new(2);
291313

292-
assert!(cache.insert_with_dirty(0, 0, false).is_none());
293-
assert!(cache.insert_with_dirty(1, 1, false).is_none());
294-
assert!(cache.insert_with_dirty(2, 2, true).is_none());
295-
assert!(cache.insert_with_dirty(3, 3, true).is_none());
314+
assert!(cache.insert_with_dirty(0, 0, false).unwrap().is_none());
315+
assert!(cache.insert_with_dirty(1, 1, false).unwrap().is_none());
316+
assert!(cache.insert_with_dirty(2, 2, true).unwrap().is_none());
317+
assert!(cache.insert_with_dirty(3, 3, true).unwrap().is_none());
296318

297319
let mut flushed = Vec::new();
298320
cache

0 commit comments

Comments
 (0)