Merge pull request #361 from cuviper/insert_entry

Add `{Entry,VacantEntry}::insert_entry`

Author: cuviper

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "indexmap"
 edition = "2021"
-version = "2.6.0"
+version = "2.7.0"
 documentation = "https://docs.rs/indexmap/"
 repository = "https://github.com/indexmap-rs/indexmap"
 license = "Apache-2.0 OR MIT"

RELEASES.md

@@ -1,5 +1,10 @@
 # Releases
 
+## 2.7.0 (2024-11-30)
+
+- Added methods `Entry::insert_entry` and `VacantEntry::insert_entry`, returning
+  an `OccupiedEntry` after insertion.
+
 ## 2.6.0 (2024-10-01)
 
 - Implemented `Clone` for `map::IntoIter` and `set::IntoIter`.

src/map/core.rs

@@ -316,6 +316,17 @@ impl<K, V> IndexMapCore<K, V> {
         self.indices.find(hash.get(), eq).copied()
     }
 
+    /// Append a key-value pair to `entries`,
+    /// *without* checking whether it already exists.
+    fn push_entry(&mut self, hash: HashValue, key: K, value: V) {
+        if self.entries.len() == self.entries.capacity() {
+            // Reserve our own capacity synced to the indices,
+            // rather than letting `Vec::push` just double it.
+            self.borrow_mut().reserve_entries(1);
+        }
+        self.entries.push(Bucket { hash, key, value });
+    }
+
     pub(crate) fn insert_full(&mut self, hash: HashValue, key: K, value: V) -> (usize, Option<V>)
     where
         K: Eq,
@@ -330,7 +341,7 @@ impl<K, V> IndexMapCore<K, V> {
             hash_table::Entry::Vacant(entry) => {
                 let i = self.entries.len();
                 entry.insert(i);
-                self.borrow_mut().push_entry(hash, key, value);
+                self.push_entry(hash, key, value);
                 debug_assert_eq!(self.indices.len(), self.entries.len());
                 (i, None)
             }
@@ -362,7 +373,7 @@ impl<K, V> IndexMapCore<K, V> {
             hash_table::Entry::Vacant(entry) => {
                 let i = self.entries.len();
                 entry.insert(i);
-                self.borrow_mut().push_entry(hash, key, value);
+                self.push_entry(hash, key, value);
                 debug_assert_eq!(self.indices.len(), self.entries.len());
                 (i, None)
             }
@@ -522,37 +533,25 @@ impl<'a, K, V> RefMut<'a, K, V> {
         self.entries.reserve_exact(additional);
     }
 
-    /// Append a key-value pair to `entries`,
+    /// Insert a key-value pair in `entries`,
     /// *without* checking whether it already exists.
-    fn push_entry(&mut self, hash: HashValue, key: K, value: V) {
+    fn insert_unique(mut self, hash: HashValue, key: K, value: V) -> OccupiedEntry<'a, K, V> {
         if self.entries.len() == self.entries.capacity() {
             // Reserve our own capacity synced to the indices,
             // rather than letting `Vec::push` just double it.
             self.reserve_entries(1);
         }
-        self.entries.push(Bucket { hash, key, value });
-    }
-
-    /// Insert a key-value pair in `entries` at a particular index,
-    /// *without* checking whether it already exists.
-    fn insert_entry(&mut self, index: usize, hash: HashValue, key: K, value: V) {
-        if self.entries.len() == self.entries.capacity() {
-            // Reserve our own capacity synced to the indices,
-            // rather than letting `Vec::insert` just double it.
-            self.reserve_entries(1);
-        }
-        self.entries.insert(index, Bucket { hash, key, value });
-    }
-
-    fn insert_unique(&mut self, hash: HashValue, key: K, value: V) -> usize {
         let i = self.indices.len();
-        self.indices
+        let entry = self
+            .indices
             .insert_unique(hash.get(), i, get_hash(self.entries));
         debug_assert_eq!(i, self.entries.len());
-        self.push_entry(hash, key, value);
-        i
+        self.entries.push(Bucket { hash, key, value });
+        OccupiedEntry::new(self.entries, entry)
     }
 
+    /// Insert a key-value pair in `entries` at a particular index,
+    /// *without* checking whether it already exists.
     fn shift_insert_unique(&mut self, index: usize, hash: HashValue, key: K, value: V) {
         let end = self.indices.len();
         assert!(index <= end);
@@ -565,7 +564,12 @@ impl<'a, K, V> RefMut<'a, K, V> {
             let i = if i < index { i } else { i - 1 };
             entries[i].hash.get()
         });
-        self.insert_entry(index, hash, key, value);
+        if self.entries.len() == self.entries.capacity() {
+            // Reserve our own capacity synced to the indices,
+            // rather than letting `Vec::insert` just double it.
+            self.reserve_entries(1);
+        }
+        self.entries.insert(index, Bucket { hash, key, value });
     }
 
     /// Remove an entry by shifting all entries that follow it

src/map/core/entry.rs

@@ -39,6 +39,19 @@ impl<'a, K, V> Entry<'a, K, V> {
         }
     }
 
+    /// Sets the value of the entry (after inserting if vacant), and returns an `OccupiedEntry`.
+    ///
+    /// Computes in **O(1)** time (amortized average).
+    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> {
+        match self {
+            Entry::Occupied(mut entry) => {
+                entry.insert(value);
+                entry
+            }
+            Entry::Vacant(entry) => entry.insert_entry(value),
+        }
+    }
+
     /// Inserts the given default value in the entry if it is vacant and returns a mutable
     /// reference to it. Otherwise a mutable reference to an already existent value is returned.
     ///
@@ -136,6 +149,13 @@ pub struct OccupiedEntry<'a, K, V> {
 }
 
 impl<'a, K, V> OccupiedEntry<'a, K, V> {
+    pub(crate) fn new(
+        entries: &'a mut Entries<K, V>,
+        index: hash_table::OccupiedEntry<'a, usize>,
+    ) -> Self {
+        Self { entries, index }
+    }
+
     /// Return the index of the key-value pair
     #[inline]
     pub fn index(&self) -> usize {
@@ -182,6 +202,11 @@ impl<'a, K, V> OccupiedEntry<'a, K, V> {
         &mut self.entries[index].value
     }
 
+    pub(super) fn into_muts(self) -> (&'a mut K, &'a mut V) {
+        let index = self.index();
+        self.entries[index].muts()
+    }
+
     /// Sets the value of the entry to `value`, and returns the entry's old value.
     pub fn insert(&mut self, value: V) -> V {
         mem::replace(self.get_mut(), value)
@@ -343,9 +368,18 @@ impl<'a, K, V> VacantEntry<'a, K, V> {
 
     /// Inserts the entry's key and the given value into the map, and returns a mutable reference
     /// to the value.
-    pub fn insert(mut self, value: V) -> &'a mut V {
-        let i = self.map.insert_unique(self.hash, self.key, value);
-        &mut self.map.entries[i].value
+    ///
+    /// Computes in **O(1)** time (amortized average).
+    pub fn insert(self, value: V) -> &'a mut V {
+        self.insert_entry(value).into_mut()
+    }
+
+    /// Inserts the entry's key and the given value into the map, and returns an `OccupiedEntry`.
+    ///
+    /// Computes in **O(1)** time (amortized average).
+    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> {
+        let Self { map, hash, key } = self;
+        map.insert_unique(hash, key, value)
     }
 
     /// Inserts the entry's key and the given value into the map at its ordered

src/map/core/raw_entry_v1.rs

@@ -618,10 +618,9 @@ impl<'a, K, V, S> RawVacantEntryMut<'a, K, V, S> {
 
     /// Inserts the given key and value into the map with the provided hash,
     /// and returns mutable references to them.
-    pub fn insert_hashed_nocheck(mut self, hash: u64, key: K, value: V) -> (&'a mut K, &'a mut V) {
+    pub fn insert_hashed_nocheck(self, hash: u64, key: K, value: V) -> (&'a mut K, &'a mut V) {
         let hash = HashValue(hash as usize);
-        let i = self.map.insert_unique(hash, key, value);
-        self.map.entries[i].muts()
+        self.map.insert_unique(hash, key, value).into_muts()
     }
 
     /// Inserts the given key and value into the map at the given index,