Add get_each_mut methods on RawTable and HashMap

These methods enable looking up mutable references to several entries in
a table or map at once. They make use of the min_const_generics feature,
which is available without a feature gate on recent nightly — but not
yet stable — rustc. Hence everything added here is behind `#[cfg(feature
= "nightly")]`.

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`.
 ///

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<Q: ?Sized, const N: usize>(
+ &mut self,
+ ks: [&Q; N],
+ ) -> [Result<&'_ mut V, UnavailableMutError>; N]
+ where
+ K: Borrow<Q>,
+ Q: Hash + Eq,
+ {
+ let mut pairs = self.get_each_inner_mut(ks);
+ let mut out: [MaybeUninit<Result<&'_ mut V, UnavailableMutError>>; 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<Q: ?Sized, const N: usize>(
+ &mut self,
+ ks: [&Q; N],
+ ) -> [Result<(&'_ K, &'_ mut V), UnavailableMutError>; N]
+ where
+ K: Borrow<Q>,
+ Q: Hash + Eq,
+ {
+ let mut pairs = self.get_each_inner_mut(ks);
+ let mut out: [MaybeUninit<Result<(&'_ K, &'_ mut V), UnavailableMutError>>; 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<Q: ?Sized, const N: usize>(
+ &mut self,
+ ks: [&Q; N],
+ ) -> [Result<&'_ mut (K, V), UnavailableMutError>; N]
+ where
+ K: Borrow<Q>,
+ Q: Hash + Eq,
+ {
+ let mut hashes = [0_u64; N];
+ for i in 0..N {
+ hashes[i] = make_hash::<K, Q, S>(&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),
+ ]
+ );
+ }
 }

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<A> {
  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<T, const N: usize>(
+ array: [MaybeUninit<T>; N],
+) -> [T; N] {
+ core::intrinsics::assert_inhabited::<[T; N]>();
+ (&array as *const _ as *const [T; N]).read()
+}
+
 impl<T> RawTable<T, Global> {
  /// Creates a new empty hash table without allocating any memory.
  ///
@@ -944,6 +962,62 @@ impl<T, A: Allocator + Clone> RawTable<T, A> {
  }
  }
 
+ /// 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<const N: usize>(
+ &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<Option<Bucket<T>>>; 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<Bucket<T>>; N] = unsafe { array_assume_init_shim(buckets) };
+
+ // Walk through the buckets, checking for duplicates and building up the output array.
+ let mut out: [MaybeUninit<Result<&'_ mut T, UnavailableMutError>>; 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