diff --git a/src/raw.rs b/src/raw.rs index 4bc1f7439..eb4d7cfa3 100644 --- a/src/raw.rs +++ b/src/raw.rs @@ -1,7 +1,7 @@ use crate::TryReserveError; use crate::control::{BitMaskIter, Group, Tag, TagSliceExt}; use crate::scopeguard::{ScopeGuard, guard}; -use crate::util::{invalid_mut, likely, unlikely}; +use crate::util::{cold_path, invalid_mut, likely, unlikely}; use core::array; use core::iter::FusedIterator; use core::marker::PhantomData; @@ -1030,32 +1030,68 @@ impl RawTable { } } - /// Inserts a new element into the table, and returns its raw bucket. + /// Finds an insert slot for the given hash if there is spare capacity. /// - /// This does not check if the given element already exists in the table. - #[cfg_attr(feature = "inline-more", inline)] - pub(crate) fn insert(&mut self, hash: u64, value: T, hasher: impl Fn(&T) -> u64) -> Bucket { + /// Returns `None` if the table is at its load-factor limit and would need + /// to grow. + #[inline] + fn find_insert_slot(&self, hash: u64) -> Option { unsafe { // SAFETY: // 1. The [`RawTableInner`] must already have properly initialized control bytes since // we will never expose `RawTable::new_uninitialized` in a public API. - // - // 2. We reserve additional space (if necessary) right after calling this function. - let mut index = self.table.find_insert_index(hash); + let index = self.table.find_insert_index(hash); // We can avoid growing the table once we have reached our load factor if we are replacing // a tombstone. This works since the number of EMPTY slots does not change in this case. // // SAFETY: The function is guaranteed to return an index in the range `0..=self.num_buckets()`. let old_ctrl = *self.table.ctrl(index); - if unlikely(self.table.growth_left == 0 && old_ctrl.special_is_empty()) { + if self.table.growth_left == 0 && old_ctrl.special_is_empty() { + None + } else { + Some(index) + } + } + } + + /// Inserts a new element into the table, and returns its raw bucket. + /// + /// This does not check if the given element already exists in the table. + #[cfg_attr(feature = "inline-more", inline)] + pub(crate) fn insert(&mut self, hash: u64, value: T, hasher: impl Fn(&T) -> u64) -> Bucket { + let index = match self.find_insert_slot(hash) { + Some(index) => index, + None => { + cold_path(); self.reserve(1, hasher); // SAFETY: We know for sure that `RawTableInner` has control bytes // initialized and that there is extra space in the table. - index = self.table.find_insert_index(hash); + unsafe { self.table.find_insert_index(hash) } } + }; + // SAFETY: `index` came from either `find_insert_slot` (which returns a + // valid `find_insert_index` result) or from a fresh `find_insert_index` + // after reserving space. + unsafe { self.insert_at_index(hash, index, value) } + } - self.insert_at_index(hash, index, value) + /// Tries to insert a new element into the table without growing. + /// + /// Returns `Ok(Bucket)` on success, or `Err(value)` if the table is at + /// its load-factor limit and would need to grow. + /// + /// This does not check if the given element already exists in the table. + #[inline] + pub(crate) fn try_insert_within_capacity( + &mut self, + hash: u64, + value: T, + ) -> Result, T> { + match self.find_insert_slot(hash) { + // SAFETY: `index` is a valid insert slot from `find_insert_slot`. + Some(index) => Ok(unsafe { self.insert_at_index(hash, index, value) }), + None => Err(value), } } diff --git a/src/table.rs b/src/table.rs index 094350442..5fdb2b506 100644 --- a/src/table.rs +++ b/src/table.rs @@ -707,6 +707,47 @@ where } } + /// Inserts an element into the `HashTable` with the given hash value, but + /// without checking whether an equivalent element already exists within + /// the table. If there is insufficient capacity, then this returns an + /// error holding the provided value. + /// + /// # Examples + /// + /// ``` + /// # #[cfg(feature = "nightly")] + /// # fn test() { + /// use hashbrown::{HashTable, DefaultHashBuilder}; + /// use std::hash::BuildHasher; + /// + /// let mut v = HashTable::new(); + /// let hasher = DefaultHashBuilder::default(); + /// let hasher = |val: &_| hasher.hash_one(val); + /// assert!(v.try_insert_unique_within_capacity(hasher(&1), 1).is_err()); + /// assert!(v.is_empty()); + /// v.reserve(1, hasher); + /// assert!(v.try_insert_unique_within_capacity(hasher(&1), 1).is_ok()); + /// assert!(!v.is_empty()); + /// # } + /// # fn main() { + /// # #[cfg(feature = "nightly")] + /// # test() + /// # } + /// ``` + pub fn try_insert_unique_within_capacity( + &mut self, + hash: u64, + value: T, + ) -> Result, T> { + match self.raw.try_insert_within_capacity(hash, value) { + Ok(bucket) => Ok(OccupiedEntry { + bucket, + table: self, + }), + Err(value) => Err(value), + } + } + /// Clears the table, removing all values. /// /// # Examples diff --git a/src/util.rs b/src/util.rs index 8667f34dc..2d16597b2 100644 --- a/src/util.rs +++ b/src/util.rs @@ -1,11 +1,12 @@ -// FIXME: Replace with `core::hint::{likely, unlikely}` once they are stable. +// FIXME: Replace with `core::hint::{likely, unlikely, cold_path}` once they are stable. #[cfg(feature = "nightly")] pub(crate) use core::intrinsics::{likely, unlikely}; -#[cfg(not(feature = "nightly"))] +/// Hints to the compiler that the code path calling this function is cold +/// (unlikely to be executed). #[inline(always)] #[cold] -fn cold_path() {} +pub(crate) fn cold_path() {} #[cfg(not(feature = "nightly"))] #[inline(always)]