Skip to content

Commit

Permalink
Specialize fold implementation of iterators
Browse files Browse the repository at this point in the history 
This provides 5-8% iteration speedups
  • Loading branch information
@a1phyr
a1phyr committed Nov 6, 2023
1 parent 778e235 commit deaf7eb
Show file tree
Hide file tree
Showing 4 changed files with 249 additions and 6 deletions.
78 changes: 78 additions & 0 deletions src/map.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2469,6 +2469,14 @@ impl<K, V, A: Allocator> Iterator for IntoKeys<K, V, A> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[inline]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, (k, _)| f(acc, k))
}
}

impl<K, V, A: Allocator> ExactSizeIterator for IntoKeys<K, V, A> {
Expand Down Expand Up @@ -2531,6 +2539,14 @@ impl<K, V, A: Allocator> Iterator for IntoValues<K, V, A> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[inline]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, (_, v)| f(acc, v))
}
}

impl<K, V, A: Allocator> ExactSizeIterator for IntoValues<K, V, A> {
Expand Down Expand Up @@ -4722,6 +4738,17 @@ impl<'a, K, V> Iterator for Iter<'a, K, V> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, x| unsafe {
let (k, v) = x.as_ref();
f(acc, (k, v))
})
}
}
impl<K, V> ExactSizeIterator for Iter<'_, K, V> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -4750,6 +4777,17 @@ impl<'a, K, V> Iterator for IterMut<'a, K, V> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, x| unsafe {
let (k, v) = x.as_mut();
f(acc, (k, v))
})
}
}
impl<K, V> ExactSizeIterator for IterMut<'_, K, V> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -4780,6 +4818,14 @@ impl<K, V, A: Allocator> Iterator for IntoIter<K, V, A> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, f)
}
}
impl<K, V, A: Allocator> ExactSizeIterator for IntoIter<K, V, A> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -4810,6 +4856,14 @@ impl<'a, K, V> Iterator for Keys<'a, K, V> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, (k, _)| f(acc, k))
}
}
impl<K, V> ExactSizeIterator for Keys<'_, K, V> {
#[cfg_attr(feature = "inline-more", inline)]
Expand All @@ -4834,6 +4888,14 @@ impl<'a, K, V> Iterator for Values<'a, K, V> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, (_, v)| f(acc, v))
}
}
impl<K, V> ExactSizeIterator for Values<'_, K, V> {
#[cfg_attr(feature = "inline-more", inline)]
Expand All @@ -4858,6 +4920,14 @@ impl<'a, K, V> Iterator for ValuesMut<'a, K, V> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, |acc, (_, v)| f(acc, v))
}
}
impl<K, V> ExactSizeIterator for ValuesMut<'_, K, V> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -4886,6 +4956,14 @@ impl<'a, K, V, A: Allocator> Iterator for Drain<'a, K, V, A> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.inner.fold(init, f)
}
}
impl<K, V, A: Allocator> ExactSizeIterator for Drain<'_, K, V, A> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down
75 changes: 71 additions & 4 deletions src/raw/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,13 +50,36 @@ use self::imp::Group;

// Branch prediction hint. This is currently only available on nightly but it
// consistently improves performance by 10-15%.
#[cfg(not(feature = "nightly"))]
use core::convert::identity as likely;
#[cfg(not(feature = "nightly"))]
use core::convert::identity as unlikely;
// #[cfg(not(feature = "nightly"))]
// use core::convert::identity as likely;
// #[cfg(not(feature = "nightly"))]
// use core::convert::identity as unlikely;
#[cfg(feature = "nightly")]
use core::intrinsics::{likely, unlikely};

#[cfg(not(feature = "nightly"))]
#[inline]
#[cold]
fn cold() {}

#[cfg(not(feature = "nightly"))]
#[inline]
fn likely(b: bool) -> bool {
if !b {
cold()
}
b
}
#[cfg(not(feature = "nightly"))]
#[cold]
#[inline]
fn unlikely(b: bool) -> bool {
if b {
cold()
}
b
}

// Use strict provenance functions if available.
#[cfg(feature = "nightly")]
use core::ptr::invalid_mut;
Expand Down Expand Up @@ -3846,6 +3869,41 @@ impl<T> RawIterRange<T> {
self.next_ctrl = self.next_ctrl.add(Group::WIDTH);
}
}

/// # Safety
/// The provided `n` value must match the actual number of items
#[allow(clippy::while_let_on_iterator)]
#[cfg_attr(feature = "inline-more", inline)]
unsafe fn fold_impl<F, B>(mut self, mut n: usize, mut acc: B, mut f: F) -> B
where
F: FnMut(B, Bucket<T>) -> B,
{
if n == 0 {
return acc;
}

loop {
while let Some(index) = self.current_group.next() {
debug_assert!(n != 0);
let bucket = self.data.next_n(index);
acc = f(acc, bucket);
n -= 1;
}

if n == 0 {
return acc;
}

// We might read past self.end up to the next group boundary,
// but this is fine because it only occurs on tables smaller
// than the group size where the trailing control bytes are all
// EMPTY. On larger tables self.end is guaranteed to be aligned
// to the group size (since tables are power-of-two sized).
self.current_group = Group::load_aligned(self.next_ctrl).match_full().into_iter();
self.data = self.data.next_n(Group::WIDTH);
self.next_ctrl = self.next_ctrl.add(Group::WIDTH);
}
}
}

// We make raw iterators unconditionally Send and Sync, and let the PhantomData
Expand Down Expand Up @@ -4069,6 +4127,15 @@ impl<T> Iterator for RawIter<T> {
fn size_hint(&self) -> (usize, Option<usize>) {
(self.items, Some(self.items))
}

#[inline]
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
unsafe { self.iter.fold_impl(self.items, init, f) }
}
}

impl<T> ExactSizeIterator for RawIter<T> {}
Expand Down
68 changes: 68 additions & 0 deletions src/set.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1696,6 +1696,14 @@ impl<'a, K> Iterator for Iter<'a, K> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, f)
}
}
impl<'a, K> ExactSizeIterator for Iter<'a, K> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -1726,6 +1734,14 @@ impl<K, A: Allocator> Iterator for IntoIter<K, A> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, |acc, (k, ())| f(acc, k))
}
}
impl<K, A: Allocator> ExactSizeIterator for IntoIter<K, A> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -1757,6 +1773,14 @@ impl<K, A: Allocator> Iterator for Drain<'_, K, A> {
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, |acc, (k, ())| f(acc, k))
}
}
impl<K, A: Allocator> ExactSizeIterator for Drain<'_, K, A> {
#[cfg_attr(feature = "inline-more", inline)]
Expand Down Expand Up @@ -1827,6 +1851,20 @@ where
let (_, upper) = self.iter.size_hint();
(0, upper)
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, |acc, elt| {
if self.other.contains(elt) {
f(acc, elt)
} else {
acc
}
})
}
}

impl<T, S, A> fmt::Debug for Intersection<'_, T, S, A>
Expand Down Expand Up @@ -1881,6 +1919,20 @@ where
let (_, upper) = self.iter.size_hint();
(0, upper)
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, mut f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, |acc, elt| {
if self.other.contains(elt) {
acc
} else {
f(acc, elt)
}
})
}
}

impl<T, S, A> FusedIterator for Difference<'_, T, S, A>
Expand Down Expand Up @@ -1927,6 +1979,14 @@ where
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, f)
}
}

impl<T, S, A> FusedIterator for SymmetricDifference<'_, T, S, A>
Expand Down Expand Up @@ -1992,6 +2052,14 @@ where
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
#[cfg_attr(feature = "inline-more", inline)]
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, f)
}
}

/// A view into a single entry in a set, which may either be vacant or occupied.
Expand Down
Loading

0 comments on commit deaf7eb

Please sign in to comment.