diff --git a/src/lib.rs b/src/lib.rs index b1c98a8046..c920af93d2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -126,6 +126,20 @@ pub enum TryReserveError { }, } +/// The error type for [`RawTable::get_each_mut`](crate::raw::RawTable::get_each_mut), +/// [`HashMap::get_each_mut`], and [`HashMap::get_each_key_value_mut`]. +#[cfg(feature = "nightly")] +#[derive(Clone, PartialEq, Eq, Debug)] +pub enum UnavailableMutError { + /// The requested entry is not present in the table. + Absent, + /// The requested entry is present, but a mutable reference to it was already created and + /// returned from this call to `get_each_mut` or `get_each_key_value_mut`. + /// + /// Includes the index of the existing mutable reference in the returned array. + Duplicate(usize), +} + /// Wrapper around `Bump` which allows it to be used as an allocator for /// `HashMap`, `HashSet` and `RawTable`. /// diff --git a/src/map.rs b/src/map.rs index 5a2eca9628..aec81a250c 100644 --- a/src/map.rs +++ b/src/map.rs @@ -1,11 +1,17 @@ +#[cfg(feature = "nightly")] +use crate::raw::array_assume_init_shim; use crate::raw::{Allocator, Bucket, Global, RawDrain, RawIntoIter, RawIter, RawTable}; use crate::TryReserveError; +#[cfg(feature = "nightly")] +use crate::UnavailableMutError; use core::borrow::Borrow; use core::fmt::{self, Debug}; use core::hash::{BuildHasher, Hash}; use core::iter::{FromIterator, FusedIterator}; use core::marker::PhantomData; use core::mem; +#[cfg(feature = "nightly")] +use core::mem::MaybeUninit; use core::ops::Index; /// Default hasher for `HashMap`. @@ -1113,6 +1119,135 @@ where self.table.get_mut(hash, equivalent_key(k)) } + /// Attempts to get mutable references to `N` values in the map at once. + /// + /// Returns an array of length `N` with the results of each query. For soundness, + /// at most one mutable reference will be returned to any value. An + /// `Err(UnavailableMutError::Duplicate(i))` in the returned array indicates that a suitable + /// key-value pair exists, but a mutable reference to the value already occurs at index `i` in + /// the returned array. + /// + /// This method is available only if the `nightly` feature is enabled. + /// + /// ``` + /// use hashbrown::{HashMap, UnavailableMutError}; + /// + /// let mut libraries = HashMap::new(); + /// libraries.insert("Bodleian Library".to_string(), 1602); + /// libraries.insert("Athenæum".to_string(), 1807); + /// libraries.insert("Herzogin-Anna-Amalia-Bibliothek".to_string(), 1691); + /// libraries.insert("Library of Congress".to_string(), 1800); + /// + /// let got = libraries.get_each_mut([ + /// "Athenæum", + /// "New York Public Library", + /// "Athenæum", + /// "Library of Congress", + /// ]); + /// assert_eq!( + /// got, + /// [ + /// Ok(&mut 1807), + /// Err(UnavailableMutError::Absent), + /// Err(UnavailableMutError::Duplicate(0)), + /// Ok(&mut 1800), + /// ] + /// ); + /// ``` + #[cfg(feature = "nightly")] + pub fn get_each_mut( + &mut self, + ks: [&Q; N], + ) -> [Result<&'_ mut V, UnavailableMutError>; N] + where + K: Borrow, + Q: Hash + Eq, + { + let mut pairs = self.get_each_inner_mut(ks); + let mut out: [MaybeUninit>; N] = + unsafe { MaybeUninit::uninit().assume_init() }; + for i in 0..N { + out[i] = MaybeUninit::new( + mem::replace(&mut pairs[i], Err(UnavailableMutError::Absent)).map(|(_, v)| v), + ); + } + unsafe { array_assume_init_shim(out) } + } + + /// Attempts to get mutable references to `N` values in the map at once, with immutable + /// references to the corresponding keys. + /// + /// Returns an array of length `N` with the results of each query. For soundness, + /// at most one mutable reference will be returned to any value. An + /// `Err(UnavailableMutError::Duplicate(i))` in the returned array indicates that a suitable + /// key-value pair exists, but a mutable reference to the value already occurs at index `i` in + /// the returned array. + /// + /// This method is available only if the `nightly` feature is enabled. + /// + /// ``` + /// use hashbrown::{HashMap, UnavailableMutError}; + /// + /// let mut libraries = HashMap::new(); + /// libraries.insert("Bodleian Library".to_string(), 1602); + /// libraries.insert("Athenæum".to_string(), 1807); + /// libraries.insert("Herzogin-Anna-Amalia-Bibliothek".to_string(), 1691); + /// libraries.insert("Library of Congress".to_string(), 1800); + /// + /// let got = libraries.get_each_key_value_mut([ + /// "Bodleian Library", + /// "Herzogin-Anna-Amalia-Bibliothek", + /// "Herzogin-Anna-Amalia-Bibliothek", + /// "Gewandhaus", + /// ]); + /// assert_eq!( + /// got, + /// [ + /// Ok((&"Bodleian Library".to_string(), &mut 1602)), + /// Ok((&"Herzogin-Anna-Amalia-Bibliothek".to_string(), &mut 1691)), + /// Err(UnavailableMutError::Duplicate(1)), + /// Err(UnavailableMutError::Absent), + /// ] + /// ); + /// ``` + #[cfg(feature = "nightly")] + pub fn get_each_key_value_mut( + &mut self, + ks: [&Q; N], + ) -> [Result<(&'_ K, &'_ mut V), UnavailableMutError>; N] + where + K: Borrow, + Q: Hash + Eq, + { + let mut pairs = self.get_each_inner_mut(ks); + let mut out: [MaybeUninit>; N] = + unsafe { MaybeUninit::uninit().assume_init() }; + for i in 0..N { + out[i] = MaybeUninit::new( + mem::replace(&mut pairs[i], Err(UnavailableMutError::Absent)) + .map(|(k, v)| (&*k, v)), + ); + } + unsafe { array_assume_init_shim(out) } + } + + #[cfg(feature = "nightly")] + fn get_each_inner_mut( + &mut self, + ks: [&Q; N], + ) -> [Result<&'_ mut (K, V), UnavailableMutError>; N] + where + K: Borrow, + Q: Hash + Eq, + { + let mut hashes = [0_u64; N]; + for i in 0..N { + hashes[i] = make_hash::(&self.hash_builder, ks[i]); + } + self.table + .get_each_mut(hashes, |i, (k, _)| ks[i].eq(k.borrow())) + } + /// Inserts a key-value pair into the map. /// /// If the map did not have this key present, [`None`] is returned. @@ -3315,6 +3450,7 @@ mod test_map { use super::{HashMap, RawEntryMut}; use crate::TryReserveError::*; use rand::{rngs::SmallRng, Rng, SeedableRng}; + use std::borrow::ToOwned; use std::cell::RefCell; use std::usize; use std::vec::Vec; @@ -4682,7 +4818,6 @@ mod test_map { #[test] fn test_const_with_hasher() { use core::hash::BuildHasher; - use std::borrow::ToOwned; use std::collections::hash_map::DefaultHasher; #[derive(Clone)] @@ -4702,4 +4837,33 @@ mod test_map { map.insert(17, "seventeen".to_owned()); assert_eq!("seventeen", map[&17]); } + + #[test] + #[cfg(feature = "nightly")] + fn test_get_each_mut() { + use crate::UnavailableMutError::*; + + let mut map = HashMap::new(); + map.insert("foo".to_owned(), 0); + map.insert("bar".to_owned(), 10); + map.insert("baz".to_owned(), 20); + map.insert("qux".to_owned(), 30); + + let xs = map.get_each_mut(["foo", "dud", "foo", "qux"]); + assert_eq!( + xs, + [Ok(&mut 0), Err(Absent), Err(Duplicate(0)), Ok(&mut 30)] + ); + + let ys = map.get_each_key_value_mut(["bar", "baz", "baz", "dip"]); + assert_eq!( + ys, + [ + Ok((&"bar".to_owned(), &mut 10)), + Ok((&"baz".to_owned(), &mut 20)), + Err(Duplicate(1)), + Err(Absent), + ] + ); + } } diff --git a/src/raw/mod.rs b/src/raw/mod.rs index 24e7f13706..871aa15be0 100644 --- a/src/raw/mod.rs +++ b/src/raw/mod.rs @@ -1,11 +1,15 @@ use crate::alloc::alloc::{handle_alloc_error, Layout}; use crate::scopeguard::guard; use crate::TryReserveError; +#[cfg(feature = "nightly")] +use crate::UnavailableMutError; use core::hint; use core::iter::FusedIterator; use core::marker::PhantomData; use core::mem; use core::mem::ManuallyDrop; +#[cfg(feature = "nightly")] +use core::mem::MaybeUninit; use core::ptr::NonNull; cfg_if! { @@ -388,6 +392,20 @@ struct RawTableInner { alloc: A, } +/// We need this workaround because `transmute` can't be used on arrays of generic length. +/// +/// Nightly has a `MaybeUninit::array_assume_init` function, but it's behind a feature gate (see +/// rust-lang/rust#80908), so for now we just open-code it here. +/// +/// TODO replace this with `MaybeUninit::array_assume_init` if/once that gets stabilized +#[cfg(feature = "nightly")] +pub(crate) unsafe fn array_assume_init_shim( + array: [MaybeUninit; N], +) -> [T; N] { + core::intrinsics::assert_inhabited::<[T; N]>(); + (&array as *const _ as *const [T; N]).read() +} + impl RawTable { /// Creates a new empty hash table without allocating any memory. /// @@ -944,6 +962,62 @@ impl RawTable { } } + /// Attempts to get mutable references to `N` entries in the table at once. + /// + /// Returns an array of length `N` with the results of each query. For soundness, + /// at most one mutable reference will be returned to any entry. An + /// `Err(UnavailableMutError::Duplicate(i))` in the returned array indicates that a suitable + /// entry exists, but a mutable reference to it already occurs at index `i` in the returned + /// array. + /// + /// The `eq` argument should be a closure such that `eq(i, k)` returns true if `k` is equal to + /// the `i`th key to be looked up. + /// + /// This method is available only if the `nightly` feature is enabled. + #[cfg(feature = "nightly")] + pub fn get_each_mut( + &mut self, + hashes: [u64; N], + mut eq: impl FnMut(usize, &T) -> bool, + ) -> [Result<&'_ mut T, UnavailableMutError>; N] { + // Collect the requested buckets. + let mut buckets: [MaybeUninit>>; N] = + unsafe { MaybeUninit::uninit().assume_init() }; + for i in 0..N { + buckets[i] = MaybeUninit::new(self.find(hashes[i], |k| eq(i, k))); + } + let buckets: [Option>; N] = unsafe { array_assume_init_shim(buckets) }; + + // Walk through the buckets, checking for duplicates and building up the output array. + let mut out: [MaybeUninit>; N] = + unsafe { MaybeUninit::uninit().assume_init() }; + for i in 0..N { + out[i] = MaybeUninit::new( + #[allow(clippy::never_loop)] + 'outer: loop { + for j in 0..i { + match (&buckets[j], &buckets[i]) { + // These two buckets are the same, and we can't safely return a second + // mutable reference to the same entry. + (Some(prev), Some(cur)) if unsafe { prev.as_ptr() == cur.as_ptr() } => { + break 'outer Err(UnavailableMutError::Duplicate(j)); + } + _ => {} + } + } + // This bucket is distinct from all previous buckets (or it doesn't exist), so + // we're clear to return the result of the lookup. + break match &buckets[i] { + None => Err(UnavailableMutError::Absent), + Some(bkt) => unsafe { Ok(bkt.as_mut()) }, + }; + }, + ) + } + + unsafe { array_assume_init_shim(out) } + } + /// Returns the number of elements the map can hold without reallocating. /// /// This number is a lower bound; the table might be able to hold