Skip to content

Commit

Permalink
Use hashbrown's new single-lookup insertion
Browse files Browse the repository at this point in the history
With `find_or_find_insert_slot` and `insert_in_slot`, we can avoid an
extra insertion lookup after checking for existence. However, that also
reserves space for a new entry *before* checking existence, so it's a
little biased towards new items. For that reason, we're not using this
for map `entry` or set `replace`, as both imply a little more weight on
the expectation that it may already exist.
  • Loading branch information
cuviper committed Jun 6, 2023
1 parent ee71507 commit c37dae6
Show file tree
Hide file tree
Showing 3 changed files with 46 additions and 24 deletions.
30 changes: 17 additions & 13 deletions src/map/core.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
}
}
}

Expand Down Expand Up @@ -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
}
}

Expand Down
28 changes: 27 additions & 1 deletion src/map/core/raw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down
12 changes: 2 additions & 10 deletions src/set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down

0 comments on commit c37dae6

Please sign in to comment.