diff --git a/utils/litemap/src/lib.rs b/utils/litemap/src/lib.rs index d0b570311c6..a8310d55cd1 100644 --- a/utils/litemap/src/lib.rs +++ b/utils/litemap/src/lib.rs @@ -66,4 +66,4 @@ pub mod store; #[cfg(any(test, feature = "testing"))] pub mod testing; -pub use map::LiteMap; +pub use map::{Entry, LiteMap, OccupiedEntry, VacantEntry}; diff --git a/utils/litemap/src/map.rs b/utils/litemap/src/map.rs index 6628c09c647..e3087bcfc03 100644 --- a/utils/litemap/src/map.rs +++ b/utils/litemap/src/map.rs @@ -9,6 +9,7 @@ use alloc::boxed::Box; use alloc::vec::Vec; use core::borrow::Borrow; use core::cmp::Ordering; +use core::fmt::Debug; use core::iter::FromIterator; use core::marker::PhantomData; use core::mem; @@ -1214,6 +1215,201 @@ impl_const_get_with_index_for_integer!(i64); impl_const_get_with_index_for_integer!(i128); impl_const_get_with_index_for_integer!(isize); +/// An entry in a `LiteMap`, which may be either occupied or vacant. +#[allow(clippy::exhaustive_enums)] +pub enum Entry<'a, K, V, S> +where + K: Ord, + S: StoreMut, +{ + Occupied(OccupiedEntry<'a, K, V, S>), + Vacant(VacantEntry<'a, K, V, S>), +} + +impl Debug for Entry<'_, K, V, S> +where + K: Ord, + S: StoreMut, +{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::Occupied(arg0) => f.debug_tuple("Occupied").field(arg0).finish(), + Self::Vacant(arg0) => f.debug_tuple("Vacant").field(arg0).finish(), + } + } +} + +/// A view into an occupied entry in a `LiteMap`. +pub struct OccupiedEntry<'a, K, V, S> +where + K: Ord, + S: StoreMut, +{ + map: &'a mut LiteMap, + index: usize, +} + +impl Debug for OccupiedEntry<'_, K, V, S> +where + K: Ord, + S: StoreMut, +{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("OccupiedEntry") + .field("index", &self.index) + .finish() + } +} + +/// A view into a vacant entry in a `LiteMap`. +pub struct VacantEntry<'a, K, V, S> +where + K: Ord, + S: StoreMut, +{ + map: &'a mut LiteMap, + key: K, + index: usize, +} + +impl Debug for VacantEntry<'_, K, V, S> +where + K: Ord, + S: StoreMut, +{ + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("VacantEntry") + .field("index", &self.index) + .finish() + } +} + +impl<'a, K, V, S> Entry<'a, K, V, S> +where + K: Ord, + S: StoreMut, +{ + /// Ensures a value is in the entry by inserting the default value if empty, + /// and returns a mutable reference to the value in the entry. + pub fn or_insert(self, default: V) -> &'a mut V { + match self { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(default), + } + } + + /// Ensures a value is in the entry by inserting the result of the default function if empty, + /// and returns a mutable reference to the value in the entry. + pub fn or_default(self) -> &'a mut V + where + V: Default, + { + self.or_insert(V::default()) + } + + /// Ensures a value is in the entry by inserting the result of the default function if empty, + /// and returns a mutable reference to the value in the entry. + pub fn or_insert_with V>(self, default: F) -> &'a mut V { + match self { + Entry::Occupied(entry) => entry.into_mut(), + Entry::Vacant(entry) => entry.insert(default()), + } + } + + /// Provides in-place mutable access to an occupied entry before any + /// potential inserts into the map. + pub fn and_modify(self, f: F) -> Self + where + F: FnOnce(&mut V), + { + match self { + Entry::Occupied(mut entry) => { + f(entry.get_mut()); + Entry::Occupied(entry) + } + Entry::Vacant(entry) => Entry::Vacant(entry), + } + } +} + +impl<'a, K, V, S> OccupiedEntry<'a, K, V, S> +where + K: Ord, + S: StoreMut, +{ + /// Gets a reference to the key in the entry. + pub fn key(&self) -> &K { + #[allow(clippy::unwrap_used)] // index is valid while we have a reference to the map + self.map.values.lm_get(self.index).unwrap().0 + } + + /// Gets a reference to the value in the entry. + pub fn get(&self) -> &V { + #[allow(clippy::unwrap_used)] // index is valid while we have a reference to the map + self.map.values.lm_get(self.index).unwrap().1 + } + + /// Gets a mutable reference to the value in the entry. + pub fn get_mut(&mut self) -> &mut V { + #[allow(clippy::unwrap_used)] // index is valid while we have a reference to the map + self.map.values.lm_get_mut(self.index).unwrap().1 + } + + /// Converts the entry into a mutable reference to the value in the entry with a lifetime bound to the map. + pub fn into_mut(self) -> &'a mut V { + #[allow(clippy::unwrap_used)] // index is valid while we have a reference to the map + self.map.values.lm_get_mut(self.index).unwrap().1 + } + + /// Sets the value of the entry, and returns the entry's old value. + pub fn insert(&mut self, value: V) -> V { + mem::replace(self.get_mut(), value) + } + + /// Takes the value out of the entry, and returns it. + pub fn remove(self) -> V { + self.map.values.lm_remove(self.index).1 + } +} + +impl<'a, K, V, S> VacantEntry<'a, K, V, S> +where + K: Ord, + S: StoreMut, +{ + /// Gets a reference to the key that would be used when inserting a value through the `VacantEntry`. + pub fn key(&self) -> &K { + &self.key + } + + /// Sets the value of the entry with the `VacantEntry`'s key, and returns a mutable reference to it. + pub fn insert(self, value: V) -> &'a mut V { + // index is valid insert index that was found via binary search + // it's valid while we have a reference to the map + self.map.values.lm_insert(self.index, self.key, value); + #[allow(clippy::unwrap_used)] // we inserted at self.index above + self.map.values.lm_get_mut(self.index).unwrap().1 + } +} + +impl LiteMap +where + K: Ord, + S: StoreMut, +{ + /// Gets the entry for the given key in the map for in-place manipulation. + pub fn entry(&mut self, key: K) -> Entry { + match self.values.lm_binary_search_by(|k| k.cmp(&key)) { + Ok(index) => Entry::Occupied(OccupiedEntry { map: self, index }), + Err(index) => Entry::Vacant(VacantEntry { + map: self, + key, + index, + }), + } + } +} + #[cfg(test)] mod test { use super::*; @@ -1239,6 +1435,7 @@ mod test { assert_eq!(expected, actual); } + fn make_13() -> LiteMap { let mut result = LiteMap::new(); result.insert(1, "one"); @@ -1324,4 +1521,122 @@ mod test { } assert!(reference.is_empty()); } + + #[test] + fn entry_insert() { + let mut map: LiteMap = LiteMap::new(); + assert!(matches!(map.entry(1), Entry::Vacant(_))); + map.entry(1).or_insert("one"); + assert!(matches!(map.entry(1), Entry::Occupied(_))); + assert_eq!(map.get(&1), Some(&"one")); + } + + #[test] + fn entry_insert_with() { + let mut map: LiteMap = LiteMap::new(); + assert!(matches!(map.entry(1), Entry::Vacant(_))); + map.entry(1).or_insert_with(|| "one"); + assert!(matches!(map.entry(1), Entry::Occupied(_))); + assert_eq!(map.get(&1), Some(&"one")); + } + + #[test] + fn entry_vacant_insert() { + let mut map: LiteMap = LiteMap::new(); + if let Entry::Vacant(entry) = map.entry(1) { + entry.insert("one"); + } + assert_eq!(map.get(&1), Some(&"one")); + } + + #[test] + fn entry_occupied_get_mut() { + let mut map: LiteMap = LiteMap::new(); + map.insert(1, "one"); + if let Entry::Occupied(mut entry) = map.entry(1) { + *entry.get_mut() = "uno"; + } + assert_eq!(map.get(&1), Some(&"uno")); + } + + #[test] + fn entry_occupied_remove() { + let mut map: LiteMap = LiteMap::new(); + map.insert(1, "one"); + if let Entry::Occupied(entry) = map.entry(1) { + entry.remove(); + } + assert_eq!(map.get(&1), None); + } + + #[test] + fn entry_occupied_key() { + let mut map: LiteMap = LiteMap::new(); + map.insert(1, "one"); + if let Entry::Occupied(entry) = map.entry(1) { + assert_eq!(entry.key(), &1); + } + } + + #[test] + fn entry_occupied_get() { + let mut map: LiteMap = LiteMap::new(); + map.insert(1, "one"); + if let Entry::Occupied(entry) = map.entry(1) { + assert_eq!(entry.get(), &"one"); + } + } + + #[test] + fn entry_occupied_insert() { + let mut map: LiteMap = LiteMap::new(); + map.insert(1, "one"); + if let Entry::Occupied(mut entry) = map.entry(1) { + assert_eq!(entry.insert("uno"), "one"); + } + assert_eq!(map.get(&1), Some(&"uno")); + } + + #[test] + fn entry_vacant_key() { + let mut map: LiteMap = LiteMap::new(); + if let Entry::Vacant(entry) = map.entry(1) { + assert_eq!(entry.key(), &1); + } + } + + #[test] + fn entry_or_insert() { + let mut map: LiteMap = LiteMap::new(); + map.entry(1).or_insert("one"); + assert_eq!(map.get(&1), Some(&"one")); + map.entry(1).or_insert("uno"); + assert_eq!(map.get(&1), Some(&"one")); + } + + #[test] + fn entry_or_insert_with() { + let mut map: LiteMap = LiteMap::new(); + map.entry(1).or_insert_with(|| "one"); + assert_eq!(map.get(&1), Some(&"one")); + map.entry(1).or_insert_with(|| "uno"); + assert_eq!(map.get(&1), Some(&"one")); + } + + #[test] + fn entry_or_default() { + let mut map: LiteMap = LiteMap::new(); + map.entry(1).or_default(); + assert_eq!(map.get(&1), Some(&String::new())); + } + + #[test] + fn entry_and_modify() { + let mut map: LiteMap = LiteMap::new(); + map.entry(1).or_insert(10); + map.entry(1).and_modify(|v| *v += 5); + assert_eq!(map.get(&1), Some(&15)); + map.entry(2).and_modify(|v| *v += 5).or_insert(20); + assert_eq!(map.get(&2), Some(&20)); + } }