Add get_mut_or_insert

src/lib.rs

@@ -190,7 +190,7 @@ pub struct LruCache<K, V, S = DefaultHasher> {
     map: HashMap<KeyRef<K>, Box<LruEntry<K, V>>, 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<K, V>,
     tail: *mut LruEntry<K, V>,
 }
@@ -355,7 +355,7 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
     }
 
     // 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<LruEntry<K, V>>) {
         if self.len() == self.cap() {
@@ -504,6 +504,61 @@ impl<K: Hash + Eq, V, S: BuildHasher> LruCache<K, V, S> {
         }
     }
 
+
+    /// 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<K, V> = &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<K, V> = &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);