From 9ff163d7086a489c717a5ad3a4216b26f3aa3c65 Mon Sep 17 00:00:00 2001 From: lufte Date: Sun, 21 Jul 2024 00:17:03 -0300 Subject: [PATCH] Create versions of get_or_insert that don't move the key Four new methods are added: * get_or_insert_ref * try_get_or_insert_ref * get_or_insert_mut_ref * try_get_or_insert_mut_ref which are analog to their existing counterparts: * get_or_insert * try_get_or_insert * get_or_insert_mut * try_get_or_insert_mut with the difference of accepting a reference to the key instead of an owned object. If the key doesn't exist in the cache and needs to be moved, only then is it cloned using to_owned() so it can be owned by the cache. This is useful when cloning the key is expensive. --- src/lib.rs | 307 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 307 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index fc392b6..be7442e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -611,6 +611,57 @@ impl LruCache { } } + /// Returns a 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 reference is returned. The value referenced by the + /// key is only cloned (using `to_owned()`) if it doesn't exist in the + /// cache. + /// + /// # Example + /// + /// ``` + /// use lru::LruCache; + /// use std::num::NonZeroUsize; + /// use std::rc::Rc; + /// + /// let key1 = Rc::new("1".to_owned()); + /// let key2 = Rc::new("2".to_owned()); + /// let mut cache = LruCache::, String>::new(NonZeroUsize::new(2).unwrap()); + /// assert_eq!(cache.get_or_insert_ref(&key1, ||"One".to_owned()), "One"); + /// assert_eq!(cache.get_or_insert_ref(&key2, ||"Two".to_owned()), "Two"); + /// assert_eq!(cache.get_or_insert_ref(&key2, ||"Not two".to_owned()), "Two"); + /// assert_eq!(cache.get_or_insert_ref(&key2, ||"Again not two".to_owned()), "Two"); + /// assert_eq!(Rc::strong_count(&key1), 2); + /// assert_eq!(Rc::strong_count(&key2), 2); // key2 was only cloned once even though we + /// // queried it 3 times + /// ``` + pub fn get_or_insert_ref<'a, Q, F>(&'a mut self, k: &'_ Q, f: F) -> &'a V + where + K: Borrow, + Q: Hash + Eq + ?Sized + alloc::borrow::ToOwned, + F: FnOnce() -> V, + { + if let Some(node) = self.map.get_mut(KeyWrapper::from_ref(k)) { + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.detach(node_ptr); + self.attach(node_ptr); + + unsafe { &*(*node_ptr).val.as_ptr() } + } else { + let v = f(); + let (_, node) = self.replace_or_create_node(k.to_owned(), v); + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.attach(node_ptr); + + let keyref = unsafe { (*node_ptr).key.as_ptr() }; + self.map.insert(KeyRef { k: keyref }, node); + unsafe { &*(*node_ptr).val.as_ptr() } + } + } + /// Returns a 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 @@ -662,6 +713,61 @@ impl LruCache { } } + /// Returns a 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 reference is returned. If `FnOnce` returns `Err`, + /// returns the `Err`. The value referenced by the key is only cloned + /// (using `to_owned()`) if it doesn't exist in the cache and `FnOnce` + /// succeeds. + /// + /// # Example + /// + /// ``` + /// use lru::LruCache; + /// use std::num::NonZeroUsize; + /// use std::rc::Rc; + /// + /// let key1 = Rc::new("1".to_owned()); + /// let key2 = Rc::new("2".to_owned()); + /// let mut cache = LruCache::, String>::new(NonZeroUsize::new(2).unwrap()); + /// let f = ||->Result {Err(())}; + /// let a = ||->Result {Ok("One".to_owned())}; + /// let b = ||->Result {Ok("Two".to_owned())}; + /// assert_eq!(cache.try_get_or_insert_ref(&key1, a), Ok(&"One".to_owned())); + /// assert_eq!(cache.try_get_or_insert_ref(&key2, f), Err(())); + /// assert_eq!(cache.try_get_or_insert_ref(&key2, b), Ok(&"Two".to_owned())); + /// assert_eq!(cache.try_get_or_insert_ref(&key2, a), Ok(&"Two".to_owned())); + /// assert_eq!(Rc::strong_count(&key1), 2); + /// assert_eq!(Rc::strong_count(&key2), 2); // key2 was only cloned once even though we + /// // queried it 3 times + /// ``` + pub fn try_get_or_insert_ref<'a, Q, F, E>(&'a mut self, k: &'_ Q, f: F) -> Result<&'a V, E> + where + K: Borrow, + Q: Hash + Eq + ?Sized + alloc::borrow::ToOwned, + F: FnOnce() -> Result, + { + if let Some(node) = self.map.get_mut(KeyWrapper::from_ref(k)) { + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.detach(node_ptr); + self.attach(node_ptr); + + unsafe { Ok(&*(*node_ptr).val.as_ptr()) } + } else { + let v = f()?; + let (_, node) = self.replace_or_create_node(k.to_owned(), v); + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.attach(node_ptr); + + let keyref = unsafe { (*node_ptr).key.as_ptr() }; + self.map.insert(KeyRef { k: keyref }, node); + Ok(unsafe { &*(*node_ptr).val.as_ptr() }) + } + } + /// 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 @@ -708,6 +814,56 @@ 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. The value referenced by the + /// key is only cloned (using `to_owned()`) if it doesn't exist in the cache. + /// + /// # Example + /// + /// ``` + /// use lru::LruCache; + /// use std::num::NonZeroUsize; + /// use std::rc::Rc; + /// + /// let key1 = Rc::new("1".to_owned()); + /// let key2 = Rc::new("2".to_owned()); + /// let mut cache = LruCache::, &'static str>::new(NonZeroUsize::new(2).unwrap()); + /// cache.get_or_insert_mut_ref(&key1, ||"One"); + /// let v = cache.get_or_insert_mut_ref(&key2, ||"Two"); + /// *v = "New two"; + /// assert_eq!(cache.get_or_insert_mut_ref(&key2, ||"Two"), &mut "New two"); + /// assert_eq!(Rc::strong_count(&key1), 2); + /// assert_eq!(Rc::strong_count(&key2), 2); // key2 was only cloned once even though we + /// // queried it 2 times + /// ``` + pub fn get_or_insert_mut_ref<'a, Q, F>(&mut self, k: &'_ Q, f: F) -> &'a mut V + where + K: Borrow, + Q: Hash + Eq + ?Sized + alloc::borrow::ToOwned, + F: FnOnce() -> V, + { + if let Some(node) = self.map.get_mut(KeyWrapper::from_ref(k)) { + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.detach(node_ptr); + self.attach(node_ptr); + + unsafe { &mut *(*node_ptr).val.as_mut_ptr() } + } else { + let v = f(); + let (_, node) = self.replace_or_create_node(k.to_owned(), v); + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.attach(node_ptr); + + let keyref = unsafe { (*node_ptr).key.as_ptr() }; + self.map.insert(KeyRef { k: keyref }, node); + unsafe { &mut *(*node_ptr).val.as_mut_ptr() } + } + } + /// 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 @@ -760,6 +916,67 @@ 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. If `FnOnce` returns `Err`, + /// returns the `Err`. The value referenced by the key is only cloned + /// (using `to_owned()`) if it doesn't exist in the cache and `FnOnce` + /// succeeds. + /// + /// # Example + /// + /// ``` + /// use lru::LruCache; + /// use std::num::NonZeroUsize; + /// use std::rc::Rc; + /// + /// let key1 = Rc::new("1".to_owned()); + /// let key2 = Rc::new("2".to_owned()); + /// let mut cache = LruCache::, String>::new(NonZeroUsize::new(2).unwrap()); + /// let f = ||->Result {Err(())}; + /// let a = ||->Result {Ok("One".to_owned())}; + /// let b = ||->Result {Ok("Two".to_owned())}; + /// assert_eq!(cache.try_get_or_insert_mut_ref(&key1, a), Ok(&mut "One".to_owned())); + /// assert_eq!(cache.try_get_or_insert_mut_ref(&key2, f), Err(())); + /// if let Ok(v) = cache.try_get_or_insert_mut_ref(&key2, b) { + /// *v = "New two".to_owned(); + /// } + /// assert_eq!(cache.try_get_or_insert_mut_ref(&key2, a), Ok(&mut "New two".to_owned())); + /// assert_eq!(Rc::strong_count(&key1), 2); + /// assert_eq!(Rc::strong_count(&key2), 2); // key2 was only cloned once even though we + /// // queried it 3 times + /// ``` + pub fn try_get_or_insert_mut_ref<'a, Q, F, E>( + &'a mut self, + k: &'_ Q, + f: F, + ) -> Result<&'a mut V, E> + where + K: Borrow, + Q: Hash + Eq + ?Sized + alloc::borrow::ToOwned, + F: FnOnce() -> Result, + { + if let Some(node) = self.map.get_mut(KeyWrapper::from_ref(k)) { + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.detach(node_ptr); + self.attach(node_ptr); + + unsafe { Ok(&mut *(*node_ptr).val.as_mut_ptr()) } + } else { + let v = f()?; + let (_, node) = self.replace_or_create_node(k.to_owned(), v); + let node_ptr: *mut LruEntry = node.as_ptr(); + + self.attach(node_ptr); + + let keyref = unsafe { (*node_ptr).key.as_ptr() }; + self.map.insert(KeyRef { k: keyref }, node); + unsafe { Ok(&mut *(*node_ptr).val.as_mut_ptr()) } + } + } + /// 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. @@ -1509,6 +1726,7 @@ mod tests { use super::LruCache; use core::{fmt::Debug, num::NonZeroUsize}; use scoped_threadpool::Pool; + use std::rc::Rc; use std::sync::atomic::{AtomicUsize, Ordering}; fn assert_opt_eq(opt: Option<&V>, v: V) { @@ -1598,6 +1816,31 @@ mod tests { assert_eq!(cache.get_or_insert("lemon", || "red"), &"orange"); } + #[test] + fn test_get_or_insert_ref() { + use alloc::borrow::ToOwned; + use alloc::string::String; + + let key1 = Rc::new("1".to_owned()); + let key2 = Rc::new("2".to_owned()); + let mut cache = LruCache::, String>::new(NonZeroUsize::new(2).unwrap()); + assert!(cache.is_empty()); + assert_eq!(cache.get_or_insert_ref(&key1, || "One".to_owned()), "One"); + assert_eq!(cache.get_or_insert_ref(&key2, || "Two".to_owned()), "Two"); + assert_eq!(cache.len(), 2); + assert!(!cache.is_empty()); + assert_eq!( + cache.get_or_insert_ref(&key2, || "Not two".to_owned()), + "Two" + ); + assert_eq!( + cache.get_or_insert_ref(&key2, || "Again not two".to_owned()), + "Two" + ); + assert_eq!(Rc::strong_count(&key1), 2); + assert_eq!(Rc::strong_count(&key2), 2); + } + #[test] fn test_try_get_or_insert() { let mut cache = LruCache::new(NonZeroUsize::new(2).unwrap()); @@ -1624,6 +1867,26 @@ mod tests { ); } + #[test] + fn test_try_get_or_insert_ref() { + use alloc::borrow::ToOwned; + use alloc::string::String; + + let key1 = Rc::new("1".to_owned()); + let key2 = Rc::new("2".to_owned()); + let mut cache = LruCache::, String>::new(NonZeroUsize::new(2).unwrap()); + let f = || -> Result { Err(()) }; + let a = || -> Result { Ok("One".to_owned()) }; + let b = || -> Result { Ok("Two".to_owned()) }; + assert_eq!(cache.try_get_or_insert_ref(&key1, a), Ok(&"One".to_owned())); + assert_eq!(cache.try_get_or_insert_ref(&key2, f), Err(())); + assert_eq!(cache.try_get_or_insert_ref(&key2, b), Ok(&"Two".to_owned())); + assert_eq!(cache.try_get_or_insert_ref(&key2, a), Ok(&"Two".to_owned())); + assert_eq!(cache.len(), 2); + assert_eq!(Rc::strong_count(&key1), 2); + assert_eq!(Rc::strong_count(&key2), 2); + } + #[test] fn test_put_and_get_or_insert_mut() { let mut cache = LruCache::new(NonZeroUsize::new(2).unwrap()); @@ -1645,6 +1908,22 @@ mod tests { assert_eq!(cache.get_or_insert_mut("lemon", || "red"), &"orange"); } + #[test] + fn test_get_or_insert_mut_ref() { + use alloc::borrow::ToOwned; + use alloc::string::String; + + let key1 = Rc::new("1".to_owned()); + let key2 = Rc::new("2".to_owned()); + let mut cache = LruCache::, &'static str>::new(NonZeroUsize::new(2).unwrap()); + assert_eq!(cache.get_or_insert_mut_ref(&key1, || "One"), &mut "One"); + let v = cache.get_or_insert_mut_ref(&key2, || "Two"); + *v = "New two"; + assert_eq!(cache.get_or_insert_mut_ref(&key2, || "Two"), &mut "New two"); + assert_eq!(Rc::strong_count(&key1), 2); + assert_eq!(Rc::strong_count(&key2), 2); + } + #[test] fn test_try_get_or_insert_mut() { let mut cache = LruCache::new(NonZeroUsize::new(2).unwrap()); @@ -1665,6 +1944,34 @@ mod tests { assert_eq!(cache.try_get_or_insert_mut(4, a), Ok(&mut "b")); } + #[test] + fn test_try_get_or_insert_mut_ref() { + use alloc::borrow::ToOwned; + use alloc::string::String; + + let key1 = Rc::new("1".to_owned()); + let key2 = Rc::new("2".to_owned()); + let mut cache = LruCache::, String>::new(NonZeroUsize::new(2).unwrap()); + let f = || -> Result { Err(()) }; + let a = || -> Result { Ok("One".to_owned()) }; + let b = || -> Result { Ok("Two".to_owned()) }; + assert_eq!( + cache.try_get_or_insert_mut_ref(&key1, a), + Ok(&mut "One".to_owned()) + ); + assert_eq!(cache.try_get_or_insert_mut_ref(&key2, f), Err(())); + if let Ok(v) = cache.try_get_or_insert_mut_ref(&key2, b) { + assert_eq!(v, &mut "Two"); + *v = "New two".to_owned(); + } + assert_eq!( + cache.try_get_or_insert_mut_ref(&key2, a), + Ok(&mut "New two".to_owned()) + ); + assert_eq!(Rc::strong_count(&key1), 2); + assert_eq!(Rc::strong_count(&key2), 2); + } + #[test] fn test_put_and_get_mut() { let mut cache = LruCache::new(NonZeroUsize::new(2).unwrap());