Optimize insertion to only use a single lookup

Author: Zoxc

src/map.rs

@@ -1787,12 +1787,21 @@ where
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn insert(&mut self, k: K, v: V) -> Option<V> {
         let hash = make_insert_hash::<K, S>(&self.hash_builder, &k);
-        if let Some((_, item)) = self.table.get_mut(hash, equivalent_key(&k)) {
-            Some(mem::replace(item, v))
-        } else {
-            self.table
-                .insert(hash, (k, v), make_hasher::<_, V, S>(&self.hash_builder));
-            None
+        self.table
+            .reserve(1, make_hasher::<_, V, S>(&self.hash_builder));
+
+        unsafe {
+            let (index, found) = self.table.find_potential(hash, equivalent_key(&k));
+
+            let bucket = self.table.bucket(index);
+
+            if found {
+                Some(mem::replace(&mut bucket.as_mut().1, v))
+            } else {
+                self.table.mark_inserted(index, hash);
+                bucket.write((k, v));
+                None
+            }
         }
     }
 

src/raw/mod.rs

@@ -663,7 +663,7 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
     /// without reallocation.
     #[cfg_attr(feature = "inline-more", inline)]
     pub fn reserve(&mut self, additional: usize, hasher: impl Fn(&T) -> u64) {
-        if additional > self.table.growth_left {
+        if unlikely(additional > self.table.growth_left) {
             // Avoid `Result::unwrap_or_else` because it bloats LLVM IR.
             if self
                 .reserve_rehash(additional, hasher, Fallibility::Infallible)
@@ -832,6 +832,22 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
         }
     }
 
+    /// Searches for an element in the table,
+    /// or a potential slot where that element could be inserted.
+    #[inline]
+    pub fn find_potential(&self, hash: u64, mut eq: impl FnMut(&T) -> bool) -> (usize, bool) {
+        self.table.find_potential_inner(hash, &mut |index| unsafe {
+            eq(self.bucket(index).as_ref())
+        })
+    }
+
+    /// Marks an element in the table as inserted.
+    #[inline]
+    pub unsafe fn mark_inserted(&mut self, index: usize, hash: u64) {
+        let old_ctrl = *self.table.ctrl(index);
+        self.table.record_item_insert_at(index, old_ctrl, hash);
+    }
+
     /// Searches for an element in the table.
     #[inline]
     pub fn find(&self, hash: u64, mut eq: impl FnMut(&T) -> bool) -> Option<Bucket<T>> {
@@ -1138,6 +1154,89 @@ impl<A: Allocator + Clone> RawTableInner<A> {
         }
     }
 
+    /// Finds the position to insert something in a group.
+    #[inline]
+    fn find_insert_slot_in_group(&self, group: &Group, probe_seq: &ProbeSeq) -> Option<usize> {
+        let bit = group.match_empty_or_deleted().lowest_set_bit();
+
+        if likely(bit.is_some()) {
+            let mut index = (probe_seq.pos + bit.unwrap()) & self.bucket_mask;
+
+            // In tables smaller than the group width, trailing control
+            // bytes outside the range of the table are filled with
+            // EMPTY entries. These will unfortunately trigger a
+            // match, but once masked may point to a full bucket that
+            // is already occupied. We detect this situation here and
+            // perform a second scan starting at the beginning of the
+            // table. This second scan is guaranteed to find an empty
+            // slot (due to the load factor) before hitting the trailing
+            // control bytes (containing EMPTY).
+            unsafe {
+                if unlikely(self.is_bucket_full(index)) {
+                    debug_assert!(self.bucket_mask < Group::WIDTH);
+                    debug_assert_ne!(probe_seq.pos, 0);
+                    index = Group::load_aligned(self.ctrl(0))
+                        .match_empty_or_deleted()
+                        .lowest_set_bit_nonzero();
+                }
+            }
+
+            Some(index)
+        } else {
+            None
+        }
+    }
+
+    /// Searches for an element in the table, or a potential slot where that element could be
+    /// inserted.
+    ///
+    /// This uses dynamic dispatch to reduce the amount of code generated, but that is
+    /// eliminated by LLVM optimizations.
+    #[inline]
+    pub fn find_potential_inner(
+        &self,
+        hash: u64,
+        eq: &mut dyn FnMut(usize) -> bool,
+    ) -> (usize, bool) {
+        let mut insert_slot = None;
+
+        let h2_hash = h2(hash);
+        let mut probe_seq = self.probe_seq(hash);
+
+        loop {
+            let group = unsafe { Group::load(self.ctrl(probe_seq.pos)) };
+
+            for bit in group.match_byte(h2_hash) {
+                let index = (probe_seq.pos + bit) & self.bucket_mask;
+
+                if likely(eq(index)) {
+                    return (index, true);
+                }
+            }
+
+            // We didn't find the element we were looking for in the group, try to get an
+            // insertion slot from the group if we don't have one yet.
+            if likely(insert_slot.is_none()) {
+                insert_slot = self.find_insert_slot_in_group(&group, &probe_seq);
+            }
+
+            // Only stop the search if the group contains at least one empty element.
+            // Otherwise, the element that we are looking for might be in a following group.
+            if likely(group.match_empty().any_bit_set()) {
+                // We must have found a insert slot by now, since the current group contains at
+                // least one. For tables smaller than the group width, there will still be an
+                // empty element in the current (and only) group due to the load factor.
+                debug_assert!(insert_slot.is_some());
+                match insert_slot {
+                    Some(insert_slot) => return (insert_slot, false),
+                    None => unsafe { hint::unreachable_unchecked() },
+                }
+            }
+
+            probe_seq.move_next(self.bucket_mask);
+        }
+    }
+
     /// Searches for an empty or deleted bucket which is suitable for inserting
     /// a new element and sets the hash for that slot.
     ///
@@ -1160,27 +1259,10 @@ impl<A: Allocator + Clone> RawTableInner<A> {
         loop {
             unsafe {
                 let group = Group::load(self.ctrl(probe_seq.pos));
-                if let Some(bit) = group.match_empty_or_deleted().lowest_set_bit() {
-                    let result = (probe_seq.pos + bit) & self.bucket_mask;
-
-                    // In tables smaller than the group width, trailing control
-                    // bytes outside the range of the table are filled with
-                    // EMPTY entries. These will unfortunately trigger a
-                    // match, but once masked may point to a full bucket that
-                    // is already occupied. We detect this situation here and
-                    // perform a second scan starting at the beginning of the
-                    // table. This second scan is guaranteed to find an empty
-                    // slot (due to the load factor) before hitting the trailing
-                    // control bytes (containing EMPTY).
-                    if unlikely(self.is_bucket_full(result)) {
-                        debug_assert!(self.bucket_mask < Group::WIDTH);
-                        debug_assert_ne!(probe_seq.pos, 0);
-                        return Group::load_aligned(self.ctrl(0))
-                            .match_empty_or_deleted()
-                            .lowest_set_bit_nonzero();
-                    }
+                let index = self.find_insert_slot_in_group(&group, &probe_seq);
 
-                    return result;
+                if likely(index.is_some()) {
+                    return index.unwrap();
                 }
             }
             probe_seq.move_next(self.bucket_mask);