Skip to content

Commit

Permalink
Add get_each_mut methods on RawTable and HashMap
Browse files Browse the repository at this point in the history
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")]`.
  • Loading branch information
cole-miller committed Feb 17, 2021
1 parent 80b2c31 commit f8b4197
Show file tree
Hide file tree
Showing 3 changed files with 243 additions and 2 deletions.
17 changes: 16 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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`.
///
Expand Down
166 changes: 165 additions & 1 deletion src/map.rs
Original file line number Diff line number Diff line change
@@ -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`.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)]
Expand All @@ -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),
]
);
}
}
62 changes: 62 additions & 0 deletions src/raw/mod.rs
Original file line number Diff line number Diff line change
@@ -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! {
Expand Down Expand Up @@ -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 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 { 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
Expand Down

0 comments on commit f8b4197

Please sign in to comment.