Merge pull request #263 from cuviper/insert_in_slot

Use hashbrown's new single-lookup insertion

Author: cuviper

src/map/core.rs

@@ -275,18 +275,14 @@ impl<K, V> IndexMapCore<K, V> {
         }
     }
 
-    /// Append a key-value pair, *without* checking whether it already exists,
-    /// and return the pair's new index.
-    fn push(&mut self, hash: HashValue, key: K, value: V) -> usize {
-        let i = self.entries.len();
-        self.indices.insert(hash.get(), i, get_hash(&self.entries));
-        if i == self.entries.capacity() {
+    /// 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.reserve_entries(1);
         }
         self.entries.push(Bucket { hash, key, value });
-        i
     }
 
     /// Return the index in `entries` where an equivalent key can be found
@@ -302,9 +298,13 @@ impl<K, V> IndexMapCore<K, V> {
     where
         K: Eq,
     {
-        match self.get_index_of(hash, &key) {
-            Some(i) => (i, Some(mem::replace(&mut self.entries[i].value, value))),
-            None => (self.push(hash, key, value), None),
+        match self.find_or_insert(hash, &key) {
+            Ok(i) => (i, Some(mem::replace(&mut self.entries[i].value, value))),
+            Err(i) => {
+                debug_assert_eq!(i, self.entries.len());
+                self.push_entry(hash, key, value);
+                (i, None)
+            }
         }
     }
 
@@ -712,14 +712,18 @@ impl<'a, K, V> VacantEntry<'a, K, V> {
 
     /// Return the index where the key-value pair will be inserted.
     pub fn index(&self) -> usize {
-        self.map.len()
+        self.map.indices.len()
     }
 
     /// Inserts the entry's key and the given value into the map, and returns a mutable reference
     /// to the value.
     pub fn insert(self, value: V) -> &'a mut V {
-        let i = self.map.push(self.hash, self.key, value);
-        &mut self.map.entries[i].value
+        let i = self.index();
+        let Self { map, hash, key } = self;
+        map.indices.insert(hash.get(), i, get_hash(&map.entries));
+        debug_assert_eq!(i, map.entries.len());
+        map.push_entry(hash, key, value);
+        &mut map.entries[i].value
     }
 }
 

src/map/core/raw.rs

@@ -2,7 +2,7 @@
 //! This module encapsulates the `unsafe` access to `hashbrown::raw::RawTable`,
 //! mostly in dealing with its bucket "pointers".
 
-use super::{equivalent, Bucket, Entry, HashValue, IndexMapCore, VacantEntry};
+use super::{equivalent, get_hash, Bucket, Entry, HashValue, IndexMapCore, VacantEntry};
 use core::fmt;
 use core::mem::replace;
 use hashbrown::raw::RawTable;
@@ -48,6 +48,32 @@ impl<K, V> IndexMapCore<K, V> {
         }
     }
 
+    /// Search for a key in the table and return `Ok(entry_index)` if found.
+    /// Otherwise, insert the key and return `Err(new_index)`.
+    ///
+    /// Note that hashbrown may resize the table to reserve space for insertion,
+    /// even before checking if it's already present, so this is somewhat biased
+    /// towards new items.
+    pub(crate) fn find_or_insert(&mut self, hash: HashValue, key: &K) -> Result<usize, usize>
+    where
+        K: Eq,
+    {
+        let hash = hash.get();
+        let eq = equivalent(key, &self.entries);
+        let hasher = get_hash(&self.entries);
+        // SAFETY: We're not mutating between find and read/insert.
+        unsafe {
+            match self.indices.find_or_find_insert_slot(hash, eq, hasher) {
+                Ok(raw_bucket) => Ok(*raw_bucket.as_ref()),
+                Err(slot) => {
+                    let index = self.indices.len();
+                    self.indices.insert_in_slot(hash, slot, index);
+                    Err(index)
+                }
+            }
+        }
+    }
+
     pub(crate) fn entry(&mut self, hash: HashValue, key: K) -> Entry<'_, K, V>
     where
         K: Eq,

src/set.rs

@@ -335,16 +335,8 @@ where
     ///
     /// Computes in **O(1)** time (amortized average).
     pub fn insert_full(&mut self, value: T) -> (usize, bool) {
-        use super::map::Entry::*;
-
-        match self.map.entry(value) {
-            Occupied(e) => (e.index(), false),
-            Vacant(e) => {
-                let index = e.index();
-                e.insert(());
-                (index, true)
-            }
-        }
+        let (index, existing) = self.map.insert_full(value, ());
+        (index, existing.is_none())
     }
 
     /// Return an iterator over the values that are in `self` but not `other`.