Auto merge of #106241 - Sp00ph:vec_deque_iter_methods, r=the8472

bors · bors · commit 4507fdaaa27e · 2023-02-18T20:12:35.000Z
Implement more methods for `vec_deque::IntoIter`

This implements a couple `Iterator` methods on `vec_deque::IntoIter` (`(try_)fold`, `(try_)rfold` `advance_(back_)by`, `next_chunk`, `count` and `last`) to allow these to be more efficient than their default implementations, also allowing many other `Iterator` methods that use these under the hood to take advantage of these manual implementations. `vec::IntoIter` has similar implementations for many of these methods. This PR does not yet implement `TrustedRandomAccess` and friends, as I'm not very familiar with the required safety guarantees.

r? `@the8472` (since you also took over my last PR)
diff --git a/library/alloc/benches/vec_deque.rs b/library/alloc/benches/vec_deque.rs
@@ -1,4 +1,8 @@
-use std::collections::VecDeque;
+use core::iter::Iterator;
+use std::{
+    collections::{vec_deque, VecDeque},
+    mem,
+};
 use test::{black_box, Bencher};
 
 #[bench]
@@ -53,6 +57,146 @@ fn bench_try_fold(b: &mut Bencher) {
     b.iter(|| black_box(ring.iter().try_fold(0, |a, b| Some(a + b))))
 }
 
+/// does the memory bookkeeping to reuse the buffer of the Vec between iterations.
+/// `setup` must not modify its argument's length or capacity. `g` must not move out of its argument.
+fn into_iter_helper<
+    T: Copy,
+    F: FnOnce(&mut VecDeque<T>),
+    G: FnOnce(&mut vec_deque::IntoIter<T>),
+>(
+    v: &mut Vec<T>,
+    setup: F,
+    g: G,
+) {
+    let ptr = v.as_mut_ptr();
+    let len = v.len();
+    // ensure that the vec is full, to make sure that any wrapping from the deque doesn't
+    // access uninitialized memory.
+    assert_eq!(v.len(), v.capacity());
+
+    let mut deque = VecDeque::from(mem::take(v));
+    setup(&mut deque);
+
+    let mut it = deque.into_iter();
+    g(&mut it);
+
+    mem::forget(it);
+
+    // SAFETY: the provided functions are not allowed to modify the allocation, so the buffer is still alive.
+    // len and capacity are accurate due to the above assertion.
+    // All the elements in the buffer are still valid, because of `T: Copy` which implies `T: !Drop`.
+    mem::forget(mem::replace(v, unsafe { Vec::from_raw_parts(ptr, len, len) }));
+}
+
+#[bench]
+fn bench_into_iter(b: &mut Bencher) {
+    let len = 1024;
+    // we reuse this allocation for every run
+    let mut vec: Vec<usize> = (0..len).collect();
+    vec.shrink_to_fit();
+
+    b.iter(|| {
+        let mut sum = 0;
+        into_iter_helper(
+            &mut vec,
+            |_| {},
+            |it| {
+                for i in it {
+                    sum += i;
+                }
+            },
+        );
+        black_box(sum);
+
+        let mut sum = 0;
+        // rotating a full deque doesn't move any memory.
+        into_iter_helper(
+            &mut vec,
+            |d| d.rotate_left(len / 2),
+            |it| {
+                for i in it {
+                    sum += i;
+                }
+            },
+        );
+        black_box(sum);
+    });
+}
+
+#[bench]
+fn bench_into_iter_fold(b: &mut Bencher) {
+    let len = 1024;
+
+    // because `fold` takes ownership of the iterator,
+    // we can't prevent it from dropping the memory,
+    // so we have to bite the bullet and reallocate
+    // for every iteration.
+    b.iter(|| {
+        let deque: VecDeque<usize> = (0..len).collect();
+        assert_eq!(deque.len(), deque.capacity());
+        let sum = deque.into_iter().fold(0, |a, b| a + b);
+        black_box(sum);
+
+        // rotating a full deque doesn't move any memory.
+        let mut deque: VecDeque<usize> = (0..len).collect();
+        assert_eq!(deque.len(), deque.capacity());
+        deque.rotate_left(len / 2);
+        let sum = deque.into_iter().fold(0, |a, b| a + b);
+        black_box(sum);
+    });
+}
+
+#[bench]
+fn bench_into_iter_try_fold(b: &mut Bencher) {
+    let len = 1024;
+    // we reuse this allocation for every run
+    let mut vec: Vec<usize> = (0..len).collect();
+    vec.shrink_to_fit();
+
+    // Iterator::any uses Iterator::try_fold under the hood
+    b.iter(|| {
+        let mut b = false;
+        into_iter_helper(&mut vec, |_| {}, |it| b = it.any(|i| i == len - 1));
+        black_box(b);
+
+        into_iter_helper(&mut vec, |d| d.rotate_left(len / 2), |it| b = it.any(|i| i == len - 1));
+        black_box(b);
+    });
+}
+
+#[bench]
+fn bench_into_iter_next_chunk(b: &mut Bencher) {
+    let len = 1024;
+    // we reuse this allocation for every run
+    let mut vec: Vec<usize> = (0..len).collect();
+    vec.shrink_to_fit();
+
+    b.iter(|| {
+        let mut buf = [0; 64];
+        into_iter_helper(
+            &mut vec,
+            |_| {},
+            |it| {
+                while let Ok(a) = it.next_chunk() {
+                    buf = a;
+                }
+            },
+        );
+        black_box(buf);
+
+        into_iter_helper(
+            &mut vec,
+            |d| d.rotate_left(len / 2),
+            |it| {
+                while let Ok(a) = it.next_chunk() {
+                    buf = a;
+                }
+            },
+        );
+        black_box(buf);
+    });
+}
+
 #[bench]
 fn bench_from_array_1000(b: &mut Bencher) {
     const N: usize = 1000;
diff --git a/library/alloc/src/collections/vec_deque/into_iter.rs b/library/alloc/src/collections/vec_deque/into_iter.rs
@@ -1,5 +1,5 @@
-use core::fmt;
 use core::iter::{FusedIterator, TrustedLen};
+use core::{array, fmt, mem::MaybeUninit, ops::Try, ptr};
 
 use crate::alloc::{Allocator, Global};
 
@@ -52,6 +52,126 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
         let len = self.inner.len();
         (len, Some(len))
     }
+
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        if self.inner.len < n {
+            let len = self.inner.len;
+            self.inner.clear();
+            Err(len)
+        } else {
+            self.inner.drain(..n);
+            Ok(())
+        }
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.inner.len
+    }
+
+    fn try_fold<B, F, R>(&mut self, mut init: B, mut f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        struct Guard<'a, T, A: Allocator> {
+            deque: &'a mut VecDeque<T, A>,
+            // `consumed <= deque.len` always holds.
+            consumed: usize,
+        }
+
+        impl<'a, T, A: Allocator> Drop for Guard<'a, T, A> {
+            fn drop(&mut self) {
+                self.deque.len -= self.consumed;
+                self.deque.head = self.deque.to_physical_idx(self.consumed);
+            }
+        }
+
+        let mut guard = Guard { deque: &mut self.inner, consumed: 0 };
+
+        let (head, tail) = guard.deque.as_slices();
+
+        init = head
+            .iter()
+            .map(|elem| {
+                guard.consumed += 1;
+                // SAFETY: Because we incremented `guard.consumed`, the
+                // deque effectively forgot the element, so we can take
+                // ownership
+                unsafe { ptr::read(elem) }
+            })
+            .try_fold(init, &mut f)?;
+
+        tail.iter()
+            .map(|elem| {
+                guard.consumed += 1;
+                // SAFETY: Same as above.
+                unsafe { ptr::read(elem) }
+            })
+            .try_fold(init, &mut f)
+    }
+
+    #[inline]
+    fn fold<B, F>(mut self, init: B, mut f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        match self.try_fold(init, |b, item| Ok::<B, !>(f(b, item))) {
+            Ok(b) => b,
+            Err(e) => match e {},
+        }
+    }
+
+    #[inline]
+    fn last(mut self) -> Option<Self::Item> {
+        self.inner.pop_back()
+    }
+
+    fn next_chunk<const N: usize>(
+        &mut self,
+    ) -> Result<[Self::Item; N], array::IntoIter<Self::Item, N>> {
+        let mut raw_arr = MaybeUninit::uninit_array();
+        let raw_arr_ptr = raw_arr.as_mut_ptr().cast();
+        let (head, tail) = self.inner.as_slices();
+
+        if head.len() >= N {
+            // SAFETY: By manually adjusting the head and length of the deque, we effectively
+            // make it forget the first `N` elements, so taking ownership of them is safe.
+            unsafe { ptr::copy_nonoverlapping(head.as_ptr(), raw_arr_ptr, N) };
+            self.inner.head = self.inner.to_physical_idx(N);
+            self.inner.len -= N;
+            // SAFETY: We initialized the entire array with items from `head`
+            return Ok(unsafe { raw_arr.transpose().assume_init() });
+        }
+
+        // SAFETY: Same argument as above.
+        unsafe { ptr::copy_nonoverlapping(head.as_ptr(), raw_arr_ptr, head.len()) };
+        let remaining = N - head.len();
+
+        if tail.len() >= remaining {
+            // SAFETY: Same argument as above.
+            unsafe {
+                ptr::copy_nonoverlapping(tail.as_ptr(), raw_arr_ptr.add(head.len()), remaining)
+            };
+            self.inner.head = self.inner.to_physical_idx(N);
+            self.inner.len -= N;
+            // SAFETY: We initialized the entire array with items from `head` and `tail`
+            Ok(unsafe { raw_arr.transpose().assume_init() })
+        } else {
+            // SAFETY: Same argument as above.
+            unsafe {
+                ptr::copy_nonoverlapping(tail.as_ptr(), raw_arr_ptr.add(head.len()), tail.len())
+            };
+            let init = head.len() + tail.len();
+            // We completely drained all the deques elements.
+            self.inner.head = 0;
+            self.inner.len = 0;
+            // SAFETY: We copied all elements from both slices to the beginning of the array, so
+            // the given range is initialized.
+            Err(unsafe { array::IntoIter::new_unchecked(raw_arr, 0..init) })
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -60,10 +180,73 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
     fn next_back(&mut self) -> Option<T> {
         self.inner.pop_back()
     }
+
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let len = self.inner.len;
+        if len >= n {
+            self.inner.truncate(len - n);
+            Ok(())
+        } else {
+            self.inner.clear();
+            Err(len)
+        }
+    }
+
+    fn try_rfold<B, F, R>(&mut self, mut init: B, mut f: F) -> R
+    where
+        F: FnMut(B, Self::Item) -> R,
+        R: Try<Output = B>,
+    {
+        struct Guard<'a, T, A: Allocator> {
+            deque: &'a mut VecDeque<T, A>,
+            // `consumed <= deque.len` always holds.
+            consumed: usize,
+        }
+
+        impl<'a, T, A: Allocator> Drop for Guard<'a, T, A> {
+            fn drop(&mut self) {
+                self.deque.len -= self.consumed;
+            }
+        }
+
+        let mut guard = Guard { deque: &mut self.inner, consumed: 0 };
+
+        let (head, tail) = guard.deque.as_slices();
+
+        init = tail
+            .iter()
+            .map(|elem| {
+                guard.consumed += 1;
+                // SAFETY: See `try_fold`'s safety comment.
+                unsafe { ptr::read(elem) }
+            })
+            .try_rfold(init, &mut f)?;
+
+        head.iter()
+            .map(|elem| {
+                guard.consumed += 1;
+                // SAFETY: Same as above.
+                unsafe { ptr::read(elem) }
+            })
+            .try_rfold(init, &mut f)
+    }
+
+    #[inline]
+    fn rfold<B, F>(mut self, init: B, mut f: F) -> B
+    where
+        F: FnMut(B, Self::Item) -> B,
+    {
+        match self.try_rfold(init, |b, item| Ok::<B, !>(f(b, item))) {
+            Ok(b) => b,
+            Err(e) => match e {},
+        }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<T, A: Allocator> ExactSizeIterator for IntoIter<T, A> {
+    #[inline]
     fn is_empty(&self) -> bool {
         self.inner.is_empty()
     }
