Add get_each_mut methods on RawTable and HashMap

src/lib.rs

@@ -20,7 +20,8 @@
         extend_one,
         allocator_api,
         slice_ptr_get,
-        nonnull_slice_from_raw_parts
+        nonnull_slice_from_raw_parts,
+        maybe_uninit_array_assume_init
     )
 )]
 #![allow(
@@ -126,6 +127,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,15 @@
 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 +1117,137 @@ 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);
+        // TODO use `MaybeUninit::uninit_array` here instead once that's stable.
+        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 { MaybeUninit::array_assume_init(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);
+        // TODO use `MaybeUninit::uninit_array` here instead once that's stable.
+        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 { MaybeUninit::array_assume_init(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! {
@@ -316,12 +320,12 @@ impl<T> Bucket<T> {
         }
     }
     #[cfg_attr(feature = "inline-more", inline)]
-    pub unsafe fn as_ptr(&self) -> *mut T {
+    pub fn as_ptr(&self) -> *mut T {
         if mem::size_of::<T>() == 0 {
             // Just return an arbitrary ZST pointer which is properly aligned
             mem::align_of::<T>() as *mut T
         } else {
-            self.ptr.as_ptr().sub(1)
+            unsafe { self.ptr.as_ptr().sub(1) }
         }
     }
     #[cfg_attr(feature = "inline-more", inline)]
@@ -944,6 +948,64 @@ 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.
+        // TODO use `MaybeUninit::uninit_array` here instead once that's stable.
+        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 { MaybeUninit::array_assume_init(buckets) };
+
+        // Walk through the buckets, checking for duplicates and building up the output array.
+        // TODO use `MaybeUninit::uninit_array` here instead once that's stable.
+        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 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 { MaybeUninit::array_assume_init(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