From ba169812e6361fee03441c1a2bcfdda74767dd3c Mon Sep 17 00:00:00 2001
From: Josh Stone <cuviper@gmail.com>
Date: Tue, 27 Aug 2024 09:34:03 -0700
Subject: [PATCH] Implement `From` between `IndexedEntry` and `OccupiedEntry`

---
 src/map/core/entry.rs | 18 ++++++++++++++++++
 src/map/core/raw.rs   | 23 +++++++++++++++++------
 src/map/tests.rs      | 25 +++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/src/map/core/entry.rs b/src/map/core/entry.rs
index ab45ecc1..5ac8c495 100644
--- a/src/map/core/entry.rs
+++ b/src/map/core/entry.rs
@@ -282,6 +282,17 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for OccupiedEntry<'_, K, V> {
     }
 }
 
+impl<'a, K, V> From<IndexedEntry<'a, K, V>> for OccupiedEntry<'a, K, V> {
+    fn from(entry: IndexedEntry<'a, K, V>) -> Self {
+        Self {
+            raw: entry
+                .map
+                .index_raw_entry(entry.index)
+                .expect("index not found"),
+        }
+    }
+}
+
 /// A view into a vacant entry in an [`IndexMap`][crate::IndexMap].
 /// It is part of the [`Entry`] enum.
 pub struct VacantEntry<'a, K, V> {
@@ -491,3 +502,10 @@ impl<K: fmt::Debug, V: fmt::Debug> fmt::Debug for IndexedEntry<'_, K, V> {
             .finish()
     }
 }
+
+impl<'a, K, V> From<OccupiedEntry<'a, K, V>> for IndexedEntry<'a, K, V> {
+    fn from(entry: OccupiedEntry<'a, K, V>) -> Self {
+        let (map, index) = entry.raw.into_inner();
+        Self { map, index }
+    }
+}
diff --git a/src/map/core/raw.rs b/src/map/core/raw.rs
index 45199433..c6a7b696 100644
--- a/src/map/core/raw.rs
+++ b/src/map/core/raw.rs
@@ -83,16 +83,19 @@ impl<K, V> IndexMapCore<K, V> {
         let entries = &*self.entries;
         let eq = move |&i: &usize| is_match(&entries[i].key);
         match self.indices.find(hash.get(), eq) {
-            // SAFETY: The entry is created with a live raw bucket, at the same time
-            // we have a &mut reference to the map, so it can not be modified further.
-            Some(raw_bucket) => Ok(RawTableEntry {
-                map: self,
-                raw_bucket,
-            }),
+            // SAFETY: The bucket is valid because we *just* found it in this map.
+            Some(raw_bucket) => Ok(unsafe { RawTableEntry::new(self, raw_bucket) }),
             None => Err(self),
         }
     }
 
+    pub(super) fn index_raw_entry(&mut self, index: usize) -> Option<RawTableEntry<'_, K, V>> {
+        let hash = self.entries.get(index)?.hash;
+        let raw_bucket = self.indices.find(hash.get(), move |&i| i == index)?;
+        // SAFETY: The bucket is valid because we *just* found it in this map.
+        Some(unsafe { RawTableEntry::new(self, raw_bucket) })
+    }
+
     pub(super) fn indices_mut(&mut self) -> impl Iterator<Item = &mut usize> {
         // SAFETY: we're not letting any of the buckets escape this function,
         // only the item references that are appropriately bound to `&mut self`.
@@ -113,6 +116,13 @@ pub(super) struct RawTableEntry<'a, K, V> {
 unsafe impl<K: Sync, V: Sync> Sync for RawTableEntry<'_, K, V> {}
 
 impl<'a, K, V> RawTableEntry<'a, K, V> {
+    /// The caller must ensure that the `raw_bucket` is valid in the given `map`,
+    /// and then we hold the `&mut` reference for exclusive access.
+    #[inline]
+    unsafe fn new(map: &'a mut IndexMapCore<K, V>, raw_bucket: RawBucket) -> Self {
+        Self { map, raw_bucket }
+    }
+
     /// Return the index of the key-value pair
     #[inline]
     pub(super) fn index(&self) -> usize {
@@ -146,6 +156,7 @@ impl<'a, K, V> RawTableEntry<'a, K, V> {
     }
 
     /// Take no action, just return the index and the original map reference.
+    #[inline]
     pub(super) fn into_inner(self) -> (&'a mut IndexMapCore<K, V>, usize) {
         let index = self.index();
         (self.map, index)
diff --git a/src/map/tests.rs b/src/map/tests.rs
index 49541181..00600892 100644
--- a/src/map/tests.rs
+++ b/src/map/tests.rs
@@ -416,6 +416,31 @@ fn get_index_entry() {
     assert_eq!(*map.get(&3).unwrap(), "4");
 }
 
+#[test]
+fn from_entries() {
+    let mut map = IndexMap::from([(1, "1"), (2, "2"), (3, "3")]);
+
+    {
+        let e = match map.entry(1) {
+            Entry::Occupied(e) => IndexedEntry::from(e),
+            Entry::Vacant(_) => panic!(),
+        };
+        assert_eq!(e.index(), 0);
+        assert_eq!(*e.key(), 1);
+        assert_eq!(*e.get(), "1");
+    }
+
+    {
+        let e = match map.get_index_entry(1) {
+            Some(e) => OccupiedEntry::from(e),
+            None => panic!(),
+        };
+        assert_eq!(e.index(), 1);
+        assert_eq!(*e.key(), 2);
+        assert_eq!(*e.get(), "2");
+    }
+}
+
 #[test]
 fn keys() {
     let vec = vec![(1, 'a'), (2, 'b'), (3, 'c')];