Implement Entry apis for LiteMap

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};

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<K, V>,
+{
+    Occupied(OccupiedEntry<'a, K, V, S>),
+    Vacant(VacantEntry<'a, K, V, S>),
+}
+
+impl<K, V, S> Debug for Entry<'_, K, V, S>
+where
+    K: Ord,
+    S: StoreMut<K, V>,
+{
+    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<K, V>,
+{
+    map: &'a mut LiteMap<K, V, S>,
+    index: usize,
+}
+
+impl<K, V, S> Debug for OccupiedEntry<'_, K, V, S>
+where
+    K: Ord,
+    S: StoreMut<K, V>,
+{
+    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<K, V>,
+{
+    map: &'a mut LiteMap<K, V, S>,
+    key: K,
+    index: usize,
+}
+
+impl<K, V, S> Debug for VacantEntry<'_, K, V, S>
+where
+    K: Ord,
+    S: StoreMut<K, V>,
+{
+    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<K, V>,
+{
+    /// 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<F: FnOnce() -> 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<F>(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<K, V>,
+{
+    /// 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<K, V>,
+{
+    /// 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<K, V, S> LiteMap<K, V, S>
+where
+    K: Ord,
+    S: StoreMut<K, V>,
+{
+    /// Gets the entry for the given key in the map for in-place manipulation.
+    pub fn entry(&mut self, key: K) -> Entry<K, V, S> {
+        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<usize, &'static str> {
         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<i32, &str> = 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<i32, &str> = 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<i32, &str> = 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<i32, &str> = 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<i32, &str> = 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<i32, &str> = 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<i32, &str> = 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<i32, &str> = 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<i32, &str> = LiteMap::new();
+        if let Entry::Vacant(entry) = map.entry(1) {
+            assert_eq!(entry.key(), &1);
+        }
+    }
+
+    #[test]
+    fn entry_or_insert() {
+        let mut map: LiteMap<i32, &str> = 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<i32, &str> = 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<i32, String> = LiteMap::new();
+        map.entry(1).or_default();
+        assert_eq!(map.get(&1), Some(&String::new()));
+    }
+
+    #[test]
+    fn entry_and_modify() {
+        let mut map: LiteMap<i32, i32> = 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));
+    }
 }