diff --git a/src/table.rs b/src/table.rs index 8b3a3fe67..b87faaad7 100644 --- a/src/table.rs +++ b/src/table.rs @@ -1071,6 +1071,24 @@ where } } + /// A version of [`iter_mut`](Self::iter_mut) which is variant over + /// elements, and thus unsafe. + /// + /// See the documentation for [`IterUnsafeMut`] for more information on how + /// to correctly use this. + /// + /// # Safety + /// + /// Any part of the returned elements which is mutated must be made invariant. + /// See the documentation for [`IterUnsafeMut`] for an example and further + /// explanation. + pub unsafe fn iter_unsafe_mut(&mut self) -> IterUnsafeMut<'_, T> { + IterUnsafeMut { + inner: unsafe { self.raw.iter() }, + marker: PhantomData, + } + } + /// An iterator producing the `usize` indices of all occupied buckets. /// /// The order in which the iterator yields indices is unspecified @@ -2504,6 +2522,16 @@ pub struct IterMut<'a, T> { inner: RawIter, marker: PhantomData<&'a mut T>, } +impl<'a, T> IterMut<'a, T> { + /// Returns a iterator of references over the remaining items. + #[cfg_attr(feature = "inline-more", inline)] + pub fn iter(&self) -> Iter<'_, T> { + Iter { + inner: self.inner.clone(), + marker: PhantomData, + } + } +} impl Default for IterMut<'_, T> { #[cfg_attr(feature = "inline-more", inline)] @@ -2552,12 +2580,113 @@ where T: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(Iter { - inner: self.inner.clone(), - marker: PhantomData, - }) - .finish() + f.debug_list().entries(self.iter()).finish() + } +} + +/// An unsafe version of [`IterMut`] which is *variant* over `T`. +/// +/// This is used for implementations of mutable iterators where mutation is only +/// allowed on part of the value in the table. For example, a mutable iterator +/// for a map may return an immutable key alongside a mutable value, even though +/// these are both stored inside the table. +/// +/// # Safety +/// +/// In order to correctly use this iterator, it should be wrapped in a safe +/// iterator struct with the appropriate [`PhantomData`] marker to indicate the +/// correct variance. +/// +/// For example, a [`hash_map::IterMut`] implementation correctly returning +/// a variant key, and an invariant value: +/// +/// [`hash_map::IterMut`]: crate::hash_map::IterMut +/// +/// ```rust +/// use core::marker::PhantomData; +/// use hashbrown::hash_table; +/// +/// pub struct IterMut<'a, K, V> { +/// inner: hash_table::IterMut<'a, (K, V)>, +/// // Variant over keys, invariant over values +/// marker: PhantomData<(&'a K, &'a mut V)>, +/// } +/// impl<'a, K, V> Iterator for IterMut<'a, K, V> { +/// // Immutable keys, mutable values +/// type Item = (&'a K, &'a mut V); +/// +/// fn next(&mut self) -> Option { +/// // Even though the key is mutably borrowed here, this is sound +/// // because we never actually mutate the key before yielding the +/// // immutable reference +/// let (ref key, ref mut val) = self.inner.next()?; +/// Some((key, val)) +/// } +/// } +/// ``` +pub struct IterUnsafeMut<'a, T> { + inner: RawIter, + marker: PhantomData<&'a T>, +} +impl<'a, T> IterUnsafeMut<'a, T> { + /// Returns a iterator of references over the remaining items. + #[cfg_attr(feature = "inline-more", inline)] + pub fn iter(&self) -> Iter<'_, T> { + Iter { + inner: self.inner.clone(), + marker: PhantomData, + } + } +} + +impl Default for IterUnsafeMut<'_, T> { + #[cfg_attr(feature = "inline-more", inline)] + fn default() -> Self { + IterUnsafeMut { + inner: Default::default(), + marker: PhantomData, + } + } +} +impl<'a, T> Iterator for IterUnsafeMut<'a, T> { + type Item = &'a mut T; + + fn next(&mut self) -> Option { + // Avoid `Option::map` because it bloats LLVM IR. + match self.inner.next() { + Some(bucket) => Some(unsafe { bucket.as_mut() }), + None => None, + } + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } + + fn fold(self, init: B, mut f: F) -> B + where + Self: Sized, + F: FnMut(B, Self::Item) -> B, + { + self.inner + .fold(init, |acc, bucket| unsafe { f(acc, bucket.as_mut()) }) + } +} + +impl ExactSizeIterator for IterUnsafeMut<'_, T> { + fn len(&self) -> usize { + self.inner.len() + } +} + +impl FusedIterator for IterUnsafeMut<'_, T> {} + +impl fmt::Debug for IterUnsafeMut<'_, T> +where + T: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list().entries(self.iter()).finish() } } @@ -2811,6 +2940,19 @@ where { inner: RawIntoIter, } +impl IntoIter +where + A: Allocator, +{ + /// Returns a iterator of references over the remaining items. + #[cfg_attr(feature = "inline-more", inline)] + pub fn iter(&self) -> Iter<'_, T> { + Iter { + inner: self.inner.iter(), + marker: PhantomData, + } + } +} impl Default for IntoIter { #[cfg_attr(feature = "inline-more", inline)] @@ -2861,12 +3003,7 @@ where A: Allocator, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(Iter { - inner: self.inner.iter(), - marker: PhantomData, - }) - .finish() + f.debug_list().entries(self.iter()).finish() } } @@ -2880,6 +3017,16 @@ where pub struct Drain<'a, T, A: Allocator = Global> { inner: RawDrain<'a, T, A>, } +impl<'a, T, A: Allocator> Drain<'a, T, A> { + /// Returns a iterator of references over the remaining items. + #[cfg_attr(feature = "inline-more", inline)] + pub fn iter(&self) -> Iter<'_, T> { + Iter { + inner: self.inner.iter(), + marker: PhantomData, + } + } +} impl Iterator for Drain<'_, T, A> { type Item = T; @@ -2911,12 +3058,7 @@ impl FusedIterator for Drain<'_, T, A> {} impl fmt::Debug for Drain<'_, T, A> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(Iter { - inner: self.inner.iter(), - marker: PhantomData, - }) - .finish() + f.debug_list().entries(self.iter()).finish() } }