Skip to content

Commit

Permalink
Alternate approach; just modify Drain
Browse files Browse the repository at this point in the history
  • Loading branch information
saethlin committed Sep 10, 2022
1 parent 25f4cb5 commit 54684c4
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 52 deletions.
44 changes: 35 additions & 9 deletions library/alloc/src/collections/vec_deque/drain.rs
Original file line number Diff line number Diff line change
@@ -1,10 +1,12 @@
use core::fmt;
use core::iter::FusedIterator;
use core::marker::PhantomData;
use core::mem::{self, MaybeUninit};
use core::ptr::{self, NonNull};
use core::{fmt, mem};

use crate::alloc::{Allocator, Global};

use super::{count, Iter, VecDeque};
use super::{count, wrap_index, VecDeque};

/// A draining iterator over the elements of a `VecDeque`.
///
Expand All @@ -20,18 +22,24 @@ pub struct Drain<
> {
after_tail: usize,
after_head: usize,
iter: Iter<'a, T>,
ring: NonNull<[T]>,
tail: usize,
head: usize,
deque: NonNull<VecDeque<T, A>>,
_phantom: PhantomData<&'a T>,
}

impl<'a, T, A: Allocator> Drain<'a, T, A> {
pub(super) unsafe fn new(
after_tail: usize,
after_head: usize,
iter: Iter<'a, T>,
ring: &'a [MaybeUninit<T>],
tail: usize,
head: usize,
deque: NonNull<VecDeque<T, A>>,
) -> Self {
Drain { after_tail, after_head, iter, deque }
let ring = unsafe { NonNull::new_unchecked(ring as *const [MaybeUninit<T>] as *mut _) };
Drain { after_tail, after_head, ring, tail, head, deque, _phantom: PhantomData }
}
}

Expand All @@ -41,7 +49,9 @@ impl<T: fmt::Debug, A: Allocator> fmt::Debug for Drain<'_, T, A> {
f.debug_tuple("Drain")
.field(&self.after_tail)
.field(&self.after_head)
.field(&self.iter)
.field(&self.ring)
.field(&self.tail)
.field(&self.head)
.finish()
}
}
Expand Down Expand Up @@ -118,20 +128,36 @@ impl<T, A: Allocator> Iterator for Drain<'_, T, A> {

#[inline]
fn next(&mut self) -> Option<T> {
self.iter.next().map(|elt| unsafe { ptr::read(elt) })
if self.tail == self.head {
return None;
}
let tail = self.tail;
self.tail = wrap_index(self.tail.wrapping_add(1), self.ring.len());
// Safety:
// - `self.tail` in a ring buffer is always a valid index.
// - `self.head` and `self.tail` equality is checked above.
unsafe { Some(ptr::read(self.ring.as_ptr().get_unchecked_mut(tail))) }
}

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
let len = count(self.tail, self.head, self.ring.len());
(len, Some(len))
}
}

#[stable(feature = "drain", since = "1.6.0")]
impl<T, A: Allocator> DoubleEndedIterator for Drain<'_, T, A> {
#[inline]
fn next_back(&mut self) -> Option<T> {
self.iter.next_back().map(|elt| unsafe { ptr::read(elt) })
if self.tail == self.head {
return None;
}
self.head = wrap_index(self.head.wrapping_sub(1), self.ring.len());
// Safety:
// - `self.head` in a ring buffer is always a valid index.
// - `self.head` and `self.tail` equality is checked above.
unsafe { Some(ptr::read(self.ring.as_ptr().get_unchecked_mut(self.head))) }
}
}

Expand Down
64 changes: 23 additions & 41 deletions library/alloc/src/collections/vec_deque/iter.rs
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
use core::fmt;
use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce};
use core::marker::PhantomData;
use core::mem::MaybeUninit;
use core::ops::Try;
use core::ptr::NonNull;

use super::{count, wrap_index, RingSlices};

Expand All @@ -15,53 +13,38 @@ use super::{count, wrap_index, RingSlices};
/// [`iter`]: super::VecDeque::iter
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Iter<'a, T: 'a> {
ring: NonNull<[T]>,
ring: &'a [MaybeUninit<T>],
tail: usize,
head: usize,
_marker: PhantomData<&'a T>,
}

impl<'a, T> Iter<'a, T> {
pub(super) fn new(ring: &'a [MaybeUninit<T>], tail: usize, head: usize) -> Self {
Iter {
ring: unsafe { NonNull::new_unchecked(ring as *const [MaybeUninit<T>] as *mut _) },
tail,
head,
_marker: PhantomData,
}
}

unsafe fn ring(&self) -> &'a [MaybeUninit<T>] {
unsafe {
core::slice::from_raw_parts(
self.ring.as_ptr() as *const MaybeUninit<T>,
self.ring.len(),
)
}
Iter { ring, tail, head }
}
}

#[stable(feature = "rust1", since = "1.0.0")]
unsafe impl<T: Sync> Sync for Iter<'_, T> {}
#[stable(feature = "rust1", since = "1.0.0")]
unsafe impl<T: Sync> Send for Iter<'_, T> {}

#[stable(feature = "collection_debug", since = "1.17.0")]
impl<T: fmt::Debug> fmt::Debug for Iter<'_, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let (front, back) = unsafe { RingSlices::ring_slices(self.ring(), self.head, self.tail) };
let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
// Safety:
// - `self.head` and `self.tail` in a ring buffer are always valid indices.
// - `RingSlices::ring_slices` guarantees that the slices split according to `self.head` and `self.tail` are initialized.
f.debug_tuple("Iter").field(&front).field(&back).finish()
unsafe {
f.debug_tuple("Iter")
.field(&MaybeUninit::slice_assume_init_ref(front))
.field(&MaybeUninit::slice_assume_init_ref(back))
.finish()
}
}
}

// FIXME(#26925) Remove in favor of `#[derive(Clone)]`
#[stable(feature = "rust1", since = "1.0.0")]
impl<T> Clone for Iter<'_, T> {
fn clone(&self) -> Self {
Iter { ring: self.ring, tail: self.tail, head: self.head, _marker: PhantomData }
Iter { ring: self.ring, tail: self.tail, head: self.head }
}
}

Expand All @@ -79,7 +62,7 @@ impl<'a, T> Iterator for Iter<'a, T> {
// Safety:
// - `self.tail` in a ring buffer is always a valid index.
// - `self.head` and `self.tail` equality is checked above.
unsafe { Some(self.ring().get_unchecked(tail).assume_init_ref()) }
unsafe { Some(self.ring.get_unchecked(tail).assume_init_ref()) }
}

#[inline]
Expand All @@ -92,11 +75,11 @@ impl<'a, T> Iterator for Iter<'a, T> {
where
F: FnMut(Acc, Self::Item) -> Acc,
{
let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
// Safety:
// - `self.head` and `self.tail` in a ring buffer are always valid indices.
// - `RingSlices::ring_slices` guarantees that the slices split according to `self.head` and `self.tail` are initialized.
unsafe {
let (front, back) = RingSlices::ring_slices(self.ring(), self.head, self.tail);
accum = MaybeUninit::slice_assume_init_ref(front).iter().fold(accum, &mut f);
MaybeUninit::slice_assume_init_ref(back).iter().fold(accum, &mut f)
}
Expand All @@ -111,13 +94,12 @@ impl<'a, T> Iterator for Iter<'a, T> {
let (mut iter, final_res);
if self.tail <= self.head {
// Safety: single slice self.ring[self.tail..self.head] is initialized.
iter =
unsafe { MaybeUninit::slice_assume_init_ref(&self.ring()[self.tail..self.head]) }
.iter();
iter = unsafe { MaybeUninit::slice_assume_init_ref(&self.ring[self.tail..self.head]) }
.iter();
final_res = iter.try_fold(init, &mut f);
} else {
// Safety: two slices: self.ring()[self.tail..], self.ring()[..self.head] both are initialized.
let (front, back) = unsafe { self.ring().split_at(self.tail) };
// Safety: two slices: self.ring[self.tail..], self.ring[..self.head] both are initialized.
let (front, back) = self.ring.split_at(self.tail);

let mut back_iter = unsafe { MaybeUninit::slice_assume_init_ref(back).iter() };
let res = back_iter.try_fold(init, &mut f);
Expand Down Expand Up @@ -151,7 +133,7 @@ impl<'a, T> Iterator for Iter<'a, T> {
// that is in bounds.
unsafe {
let idx = wrap_index(self.tail.wrapping_add(idx), self.ring.len());
self.ring().get_unchecked(idx).assume_init_ref()
self.ring.get_unchecked(idx).assume_init_ref()
}
}
}
Expand All @@ -167,18 +149,18 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
// Safety:
// - `self.head` in a ring buffer is always a valid index.
// - `self.head` and `self.tail` equality is checked above.
unsafe { Some(self.ring().get_unchecked(self.head).assume_init_ref()) }
unsafe { Some(self.ring.get_unchecked(self.head).assume_init_ref()) }
}

fn rfold<Acc, F>(self, mut accum: Acc, mut f: F) -> Acc
where
F: FnMut(Acc, Self::Item) -> Acc,
{
let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
// Safety:
// - `self.head` and `self.tail` in a ring buffer are always valid indices.
// - `RingSlices::ring_slices` guarantees that the slices split according to `self.head` and `self.tail` are initialized.
unsafe {
let (front, back) = RingSlices::ring_slices(self.ring(), self.head, self.tail);
accum = MaybeUninit::slice_assume_init_ref(back).iter().rfold(accum, &mut f);
MaybeUninit::slice_assume_init_ref(front).iter().rfold(accum, &mut f)
}
Expand All @@ -192,14 +174,14 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
{
let (mut iter, final_res);
if self.tail <= self.head {
// Safety: single slice self.ring()[self.tail..self.head] is initialized.
// Safety: single slice self.ring[self.tail..self.head] is initialized.
iter = unsafe {
MaybeUninit::slice_assume_init_ref(&self.ring()[self.tail..self.head]).iter()
MaybeUninit::slice_assume_init_ref(&self.ring[self.tail..self.head]).iter()
};
final_res = iter.try_rfold(init, &mut f);
} else {
// Safety: two slices: self.ring()[self.tail..], self.ring()[..self.head] both are initialized.
let (front, back) = unsafe { self.ring().split_at(self.tail) };
// Safety: two slices: self.ring[self.tail..], self.ring[..self.head] both are initialized.
let (front, back) = self.ring.split_at(self.tail);

let mut front_iter =
unsafe { MaybeUninit::slice_assume_init_ref(&front[..self.head]).iter() };
Expand Down
3 changes: 1 addition & 2 deletions library/alloc/src/collections/vec_deque/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1334,9 +1334,8 @@ impl<T, A: Allocator> VecDeque<T, A> {
// it. We do not write to `self` nor reborrow to a mutable reference.
// Hence the raw pointer we created above, for `deque`, remains valid.
let ring = self.buffer_as_slice();
let iter = Iter::new(ring, drain_tail, drain_head);

Drain::new(drain_head, head, iter, deque)
Drain::new(drain_head, head, ring, drain_tail, drain_head, deque)
}
}

Expand Down

0 comments on commit 54684c4

Please sign in to comment.