From 10f870eff44f07e3e8dc241e5cc11e0fa4d50ae3 Mon Sep 17 00:00:00 2001
From: oxalica <oxalicc@pm.me>
Date: Wed, 3 Aug 2022 18:16:05 +0800
Subject: [PATCH] Impl entry-API for la_arena::ArenaMap

We enforce integral and `Copy` key, so some key-related functions are
not necessary since user can just reuse the index for the `entry` call.
---
 lib/la-arena/src/lib.rs |   2 +-
 lib/la-arena/src/map.rs | 107 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 1 deletion(-)

diff --git a/lib/la-arena/src/lib.rs b/lib/la-arena/src/lib.rs
index a6766edd41..a3fe59e946 100644
--- a/lib/la-arena/src/lib.rs
+++ b/lib/la-arena/src/lib.rs
@@ -12,7 +12,7 @@ use std::{
 };
 
 mod map;
-pub use map::ArenaMap;
+pub use map::{ArenaMap, Entry, OccupiedEntry, VacantEntry};
 
 /// The raw index of a value in an arena.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
diff --git a/lib/la-arena/src/map.rs b/lib/la-arena/src/map.rs
index b5c49b3435..f74ec5d97c 100644
--- a/lib/la-arena/src/map.rs
+++ b/lib/la-arena/src/map.rs
@@ -56,6 +56,16 @@ impl<T, V> ArenaMap<Idx<T>, V> {
         self.v.iter().enumerate().filter_map(|(idx, o)| Some((Self::from_idx(idx), o.as_ref()?)))
     }
 
+    /// Gets the given key's corresponding entry in the map for in-place manipulation.
+    pub fn entry(&mut self, idx: Idx<T>) -> Entry<'_, Idx<T>, V> {
+        let idx = Self::to_idx(idx);
+        self.v.resize_with((idx + 1).max(self.v.len()), || None);
+        match &mut self.v[idx] {
+            slot @ Some(_) => Entry::Occupied(OccupiedEntry { slot, _ty: PhantomData }),
+            slot @ None => Entry::Vacant(VacantEntry { slot, _ty: PhantomData }),
+        }
+    }
+
     fn to_idx(idx: Idx<T>) -> usize {
         u32::from(idx.into_raw()) as usize
     }
@@ -83,3 +93,100 @@ impl<T, V> Default for ArenaMap<Idx<V>, T> {
         Self::new()
     }
 }
+
+/// A view into a single entry in a map, which may either be vacant or occupied.
+///
+/// This `enum` is constructed from the [`entry`] method on [`ArenaMap`].
+///
+/// [`entry`]: ArenaMap::entry
+pub enum Entry<'a, IDX, V> {
+    /// A vacant entry.
+    Vacant(VacantEntry<'a, IDX, V>),
+    /// An occupied entry.
+    Occupied(OccupiedEntry<'a, IDX, V>),
+}
+
+impl<'a, IDX, V> Entry<'a, IDX, V> {
+    /// Ensures a value is in the entry by inserting the default 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 {
+            Self::Vacant(ent) => ent.insert(default),
+            Self::Occupied(ent) => ent.into_mut(),
+        }
+    }
+
+    /// 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 {
+            Self::Vacant(ent) => ent.insert(default()),
+            Self::Occupied(ent) => ent.into_mut(),
+        }
+    }
+
+    /// Provides in-place mutable access to an occupied entry before any potential inserts into the map.
+    pub fn and_modify<F: FnOnce(&mut V)>(mut self, f: F) -> Self {
+        if let Self::Occupied(ent) = &mut self {
+            f(ent.get_mut());
+        }
+        self
+    }
+}
+
+impl<'a, IDX, V> Entry<'a, IDX, V>
+where
+    V: Default,
+{
+    /// 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_default(self) -> &'a mut V {
+        self.or_insert_with(Default::default)
+    }
+}
+
+/// A view into an vacant entry in a [`ArenaMap`]. It is part of the [`Entry`] enum.
+pub struct VacantEntry<'a, IDX, V> {
+    slot: &'a mut Option<V>,
+    _ty: PhantomData<IDX>,
+}
+
+impl<'a, IDX, V> VacantEntry<'a, IDX, V> {
+    /// 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 {
+        self.slot.insert(value)
+    }
+}
+
+/// A view into an occupied entry in a [`ArenaMap`]. It is part of the [`Entry`] enum.
+pub struct OccupiedEntry<'a, IDX, V> {
+    slot: &'a mut Option<V>,
+    _ty: PhantomData<IDX>,
+}
+
+impl<'a, IDX, V> OccupiedEntry<'a, IDX, V> {
+    /// Gets a reference to the value in the entry.
+    pub fn get(&self) -> &V {
+        self.slot.as_ref().expect("Occupied")
+    }
+
+    /// Gets a mutable reference to the value in the entry.
+    pub fn get_mut(&mut self) -> &mut V {
+        self.slot.as_mut().expect("Occupied")
+    }
+
+    /// Converts the entry into a mutable reference to its value.
+    pub fn into_mut(self) -> &'a mut V {
+        self.slot.as_mut().expect("Occupied")
+    }
+
+    /// Sets the value of the entry with the `OccupiedEntry`’s key, and returns the entry’s old value.
+    pub fn insert(&mut self, value: V) -> V {
+        self.slot.replace(value).expect("Occupied")
+    }
+
+    /// Takes the value of the entry out of the map, and returns it.
+    pub fn remove(self) -> V {
+        self.slot.take().expect("Occupied")
+    }
+}