diff --git a/src/lib.rs b/src/lib.rs index 7f09dbe..d8df14d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -190,7 +190,7 @@ pub struct LruCache { map: HashMap, Box>, S>, cap: usize, - // head and tail are sigil nodes to faciliate inserting entries + // head and tail are sigil nodes to facilitate inserting entries head: *mut LruEntry, tail: *mut LruEntry, } @@ -355,7 +355,7 @@ impl LruCache { } // Used internally to swap out a node if the cache is full or to create a new node if space - // is available. Shared between `put`, `push`, and `get_or_insert`. + // is available. Shared between `put`, `push`, `get_or_insert`, and `get_mut_or_insert`. #[allow(clippy::type_complexity)] fn replace_or_create_node(&mut self, k: K, v: V) -> (Option<(K, V)>, Box>) { if self.len() == self.cap() { @@ -504,6 +504,61 @@ impl LruCache { } } + + /// Returns a mutable reference to the value of the key in the cache if it is + /// present in the cache and moves the key to the head of the LRU list. + /// If the key does not exist the provided `FnOnce` is used to populate + /// the list and a mutable reference is returned. + /// + /// This method will only return `None` when the capacity of the cache is 0 and no entries + /// can be populated. + /// + /// # Example + /// + /// ``` + /// use lru::LruCache; + /// let mut cache = LruCache::new(2); + /// + /// cache.put(1, "a"); + /// cache.put(2, "b"); + /// + /// let v = cache.get_mut_or_insert(2, ||"c").unwrap(); + /// assert_eq!(v, &"b"); + /// *v = "d"; + /// assert_eq!(cache.get_mut_or_insert(2, ||"e"), Some(&mut "d")); + /// assert_eq!(cache.get_mut_or_insert(3, ||"f"), Some(&mut "f")); + /// assert_eq!(cache.get_mut_or_insert(3, ||"e"), Some(&mut "f")); + /// ``` + pub fn get_mut_or_insert<'a, F>(&mut self, k: K, f: F) -> Option<&'a mut V> + where + F: FnOnce() -> V, + { + if let Some(node) = self.map.get_mut(&KeyRef { k: &k }) { + let node_ptr: *mut LruEntry = &mut **node; + + self.detach(node_ptr); + self.attach(node_ptr); + + Some(unsafe { &mut (*(*node_ptr).val.as_mut_ptr()) as &mut V }) + } else { + // If the capacity is 0 we do nothing, + // this is the only circumstance that should return None + if self.cap() == 0 { + return None; + } + let v = f(); + let (_, mut node) = self.replace_or_create_node(k, v); + + let node_ptr: *mut LruEntry = &mut *node; + self.attach(node_ptr); + + let keyref = unsafe { (*node_ptr).key.as_ptr() }; + self.map.insert(KeyRef { k: keyref }, node); + Some(unsafe { &mut (*(*node_ptr).val.as_mut_ptr()) as &mut V }) + } + } + + /// Returns a reference to the value corresponding to the key in the cache or `None` if it is /// not present in the cache. Unlike `get`, `peek` does not update the LRU list so the key's /// position will be unchanged. @@ -1236,6 +1291,28 @@ mod tests { assert_opt_eq(cache.get_or_insert(&"lemon", || "red"), &"orange"); } + #[test] + fn test_put_and_get_mut_or_insert() { + let mut cache = LruCache::new(2); + assert!(cache.is_empty()); + + assert_eq!(cache.put("apple", "red"), None); + assert_eq!(cache.put("banana", "yellow"), None); + + assert_eq!(cache.cap(), 2); + assert_eq!(cache.len(), 2); + + let v = cache.get_mut_or_insert(&"apple", || "orange").unwrap(); + assert_eq!(v, &"red"); + *v = "blue"; + + assert_opt_eq_mut(cache.get_mut_or_insert(&"apple", || "orange"), &"blue"); + assert_opt_eq_mut(cache.get_mut_or_insert(&"banana", || "orange"), &"yellow"); + assert_opt_eq_mut(cache.get_mut_or_insert(&"lemon", || "orange"), &"orange"); + assert_opt_eq_mut(cache.get_mut_or_insert(&"lemon", || "red"), &"orange"); + } + + #[test] fn test_put_and_get_mut() { let mut cache = LruCache::new(2);