Skip to content

Commit

Permalink
Use MaybeUninit in VecDeque to remove the undefined behavior of slice
Browse files Browse the repository at this point in the history
Signed-off-by: JmPotato <[email protected]>
  • Loading branch information
JmPotato committed Mar 9, 2022
1 parent 10dccdc commit e54d4ab
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 29 deletions.
59 changes: 41 additions & 18 deletions library/alloc/src/collections/vec_deque/iter.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use core::fmt;
use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess, TrustedRandomAccessNoCoerce};
use core::mem::MaybeUninit;
use core::ops::Try;

use super::{count, wrap_index, RingSlices};
Expand All @@ -12,7 +13,7 @@ use super::{count, wrap_index, RingSlices};
/// [`iter`]: super::VecDeque::iter
#[stable(feature = "rust1", since = "1.0.0")]
pub struct Iter<'a, T: 'a> {
pub(crate) ring: &'a [T],
pub(crate) ring: &'a [MaybeUninit<T>],
pub(crate) tail: usize,
pub(crate) head: usize,
}
Expand Down Expand Up @@ -44,7 +45,10 @@ impl<'a, T> Iterator for Iter<'a, T> {
}
let tail = self.tail;
self.tail = wrap_index(self.tail.wrapping_add(1), self.ring.len());
unsafe { Some(self.ring.get_unchecked(tail)) }
// 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()) }
}

#[inline]
Expand All @@ -58,8 +62,13 @@ impl<'a, T> Iterator for Iter<'a, T> {
F: FnMut(Acc, Self::Item) -> Acc,
{
let (front, back) = RingSlices::ring_slices(self.ring, self.head, self.tail);
accum = front.iter().fold(accum, &mut f);
back.iter().fold(accum, &mut f)
// 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 {
accum = MaybeUninit::slice_assume_init_ref(front).iter().fold(accum, &mut f);
MaybeUninit::slice_assume_init_ref(back).iter().fold(accum, &mut f)
}
}

fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R
Expand All @@ -70,17 +79,19 @@ impl<'a, T> Iterator for Iter<'a, T> {
{
let (mut iter, final_res);
if self.tail <= self.head {
// single slice self.ring[self.tail..self.head]
iter = self.ring[self.tail..self.head].iter();
// 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();
final_res = iter.try_fold(init, &mut f);
} else {
// two slices: self.ring[self.tail..], self.ring[..self.head]
// 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 = back.iter();

let mut back_iter = unsafe { MaybeUninit::slice_assume_init_ref(back).iter() };
let res = back_iter.try_fold(init, &mut f);
let len = self.ring.len();
self.tail = (self.ring.len() - back_iter.len()) & (len - 1);
iter = front[..self.head].iter();
iter = unsafe { MaybeUninit::slice_assume_init_ref(&front[..self.head]).iter() };
final_res = iter.try_fold(res?, &mut f);
}
self.tail = self.head - iter.len();
Expand Down Expand Up @@ -109,7 +120,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)
self.ring.get_unchecked(idx).assume_init_ref()
}
}
}
Expand All @@ -122,16 +133,24 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
return None;
}
self.head = wrap_index(self.head.wrapping_sub(1), self.ring.len());
unsafe { Some(self.ring.get_unchecked(self.head)) }
// 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()) }
}

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);
accum = back.iter().rfold(accum, &mut f);
front.iter().rfold(accum, &mut f)
// 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 {
accum = MaybeUninit::slice_assume_init_ref(back).iter().rfold(accum, &mut f);
MaybeUninit::slice_assume_init_ref(front).iter().rfold(accum, &mut f)
}
}

fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R
Expand All @@ -142,16 +161,20 @@ impl<'a, T> DoubleEndedIterator for Iter<'a, T> {
{
let (mut iter, final_res);
if self.tail <= self.head {
// single slice self.ring[self.tail..self.head]
iter = self.ring[self.tail..self.head].iter();
// 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()
};
final_res = iter.try_rfold(init, &mut f);
} else {
// two slices: self.ring[self.tail..], self.ring[..self.head]
// 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 = front[..self.head].iter();

let mut front_iter =
unsafe { MaybeUninit::slice_assume_init_ref(&front[..self.head]).iter() };
let res = front_iter.try_rfold(init, &mut f);
self.head = front_iter.len();
iter = back.iter();
iter = unsafe { MaybeUninit::slice_assume_init_ref(back).iter() };
final_res = iter.try_rfold(res?, &mut f);
}
self.head = self.tail + iter.len();
Expand Down
56 changes: 45 additions & 11 deletions library/alloc/src/collections/vec_deque/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ use core::fmt;
use core::hash::{Hash, Hasher};
use core::iter::{repeat_with, FromIterator};
use core::marker::PhantomData;
use core::mem::{self, ManuallyDrop};
use core::mem::{self, ManuallyDrop, MaybeUninit};
use core::ops::{Index, IndexMut, Range, RangeBounds};
use core::ptr::{self, NonNull};
use core::slice;
Expand Down Expand Up @@ -181,16 +181,28 @@ impl<T, A: Allocator> VecDeque<T, A> {
}
}

/// Turn ptr into a slice
/// Turn ptr into a slice, since the elements of the backing buffer may be uninitialized,
/// we will return a slice of [`MaybeUninit<T>`].
///
/// See [`MaybeUninit::zeroed`][zeroed] for examples of correct and
/// incorrect usage of this method.
///
/// [zeroed]: mem::MaybeUninit::zeroed
#[inline]
unsafe fn buffer_as_slice(&self) -> &[T] {
unsafe { slice::from_raw_parts(self.ptr(), self.cap()) }
unsafe fn buffer_as_slice(&self) -> &[MaybeUninit<T>] {
unsafe { slice::from_raw_parts(self.ptr() as *mut MaybeUninit<T>, self.cap()) }
}

/// Turn ptr into a mut slice
/// Turn ptr into a mut slice, since the elements of the backing buffer may be uninitialized,
/// we will return a slice of [`MaybeUninit<T>`].
///
/// See [`MaybeUninit::zeroed`][zeroed] for examples of correct and
/// incorrect usage of this method.
///
/// [zeroed]: mem::MaybeUninit::zeroed
#[inline]
unsafe fn buffer_as_mut_slice(&mut self) -> &mut [T] {
unsafe { slice::from_raw_parts_mut(self.ptr(), self.cap()) }
unsafe fn buffer_as_mut_slice(&mut self) -> &mut [MaybeUninit<T>] {
unsafe { slice::from_raw_parts_mut(self.ptr() as *mut MaybeUninit<T>, self.cap()) }
}

/// Moves an element out of the buffer
Expand Down Expand Up @@ -1055,9 +1067,13 @@ impl<T, A: Allocator> VecDeque<T, A> {
#[inline]
#[stable(feature = "deque_extras_15", since = "1.5.0")]
pub fn as_slices(&self) -> (&[T], &[T]) {
// 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 buf = self.buffer_as_slice();
RingSlices::ring_slices(buf, self.head, self.tail)
let (front, back) = RingSlices::ring_slices(buf, self.head, self.tail);
(MaybeUninit::slice_assume_init_ref(front), MaybeUninit::slice_assume_init_ref(back))
}
}

Expand Down Expand Up @@ -1089,11 +1105,15 @@ impl<T, A: Allocator> VecDeque<T, A> {
#[inline]
#[stable(feature = "deque_extras_15", since = "1.5.0")]
pub fn as_mut_slices(&mut self) -> (&mut [T], &mut [T]) {
// 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 head = self.head;
let tail = self.tail;
let buf = self.buffer_as_mut_slice();
RingSlices::ring_slices(buf, head, tail)
let (front, back) = RingSlices::ring_slices(buf, head, tail);
(MaybeUninit::slice_assume_init_mut(front), MaybeUninit::slice_assume_init_mut(back))
}
}

Expand Down Expand Up @@ -2327,7 +2347,14 @@ impl<T, A: Allocator> VecDeque<T, A> {
if self.is_contiguous() {
let tail = self.tail;
let head = self.head;
return unsafe { RingSlices::ring_slices(self.buffer_as_mut_slice(), head, tail).0 };
// 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.
return unsafe {
MaybeUninit::slice_assume_init_mut(
RingSlices::ring_slices(self.buffer_as_mut_slice(), head, tail).0,
)
};
}

let buf = self.buf.ptr();
Expand Down Expand Up @@ -2413,7 +2440,14 @@ impl<T, A: Allocator> VecDeque<T, A> {

let tail = self.tail;
let head = self.head;
unsafe { RingSlices::ring_slices(self.buffer_as_mut_slice(), head, tail).0 }
// 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 {
MaybeUninit::slice_assume_init_mut(
RingSlices::ring_slices(self.buffer_as_mut_slice(), head, tail).0,
)
}
}

/// Rotates the double-ended queue `mid` places to the left.
Expand Down

0 comments on commit e54d4ab

Please sign in to comment.