diff --git a/src/map/core.rs b/src/map/core.rs index 908d1e50..c918deac 100644 --- a/src/map/core.rs +++ b/src/map/core.rs @@ -275,18 +275,14 @@ impl IndexMapCore { } } - /// 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 IndexMapCore { 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 } } diff --git a/src/map/core/raw.rs b/src/map/core/raw.rs index dee6363d..332770cc 100644 --- a/src/map/core/raw.rs +++ b/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 IndexMapCore { } } + /// 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 + 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, diff --git a/src/set.rs b/src/set.rs index 035d5d73..8e107594 100644 --- a/src/set.rs +++ b/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`.