Skip to content

Commit

Permalink
RepeatN: use MaybeUninit
Browse files Browse the repository at this point in the history
  • Loading branch information
fee1-dead committed Sep 9, 2024
1 parent 263a3ae commit 361c3d3
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 15 deletions.
60 changes: 45 additions & 15 deletions library/core/src/iter/sources/repeat_n.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use crate::fmt;
use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
use crate::mem::ManuallyDrop;
use crate::mem::{self, MaybeUninit};
use crate::num::NonZero;

/// Creates a new iterator that repeats a single element a given number of times.
Expand Down Expand Up @@ -58,14 +59,12 @@ use crate::num::NonZero;
#[inline]
#[stable(feature = "iter_repeat_n", since = "1.82.0")]
pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
let mut element = ManuallyDrop::new(element);

if count == 0 {
// SAFETY: we definitely haven't dropped it yet, since we only just got
// passed it in, and because the count is zero the instance we're about
// to create won't drop it, so to avoid leaking we need to now.
unsafe { ManuallyDrop::drop(&mut element) };
}
let element = if count == 0 {
// `element` gets dropped eagerly.
MaybeUninit::uninit()
} else {
MaybeUninit::new(element)
};

RepeatN { element, count }
}
Expand All @@ -74,31 +73,60 @@ pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
///
/// This `struct` is created by the [`repeat_n()`] function.
/// See its documentation for more.
#[derive(Clone, Debug)]
#[stable(feature = "iter_repeat_n", since = "1.82.0")]
pub struct RepeatN<A> {
count: usize,
// Invariant: has been dropped iff count == 0.
element: ManuallyDrop<A>,
// Invariant: uninit iff count == 0.
element: MaybeUninit<A>,
}

impl<A> RepeatN<A> {
/// Returns the element if it hasn't been dropped already.
fn element_ref(&self) -> Option<&A> {
if self.count > 0 {
// SAFETY: The count is non-zero, so it must be initialized.
Some(unsafe { self.element.assume_init_ref() })
} else {
None
}
}
/// If we haven't already dropped the element, return it in an option.
///
/// Clears the count so it won't be dropped again later.
#[inline]
fn take_element(&mut self) -> Option<A> {
if self.count > 0 {
self.count = 0;
let element = mem::replace(&mut self.element, MaybeUninit::uninit());
// SAFETY: We just set count to zero so it won't be dropped again,
// and it used to be non-zero so it hasn't already been dropped.
unsafe { Some(ManuallyDrop::take(&mut self.element)) }
unsafe { Some(element.assume_init()) }
} else {
None
}
}
}

#[stable(feature = "iter_repeat_n", since = "1.82.0")]
impl<A: Clone> Clone for RepeatN<A> {
fn clone(&self) -> RepeatN<A> {
RepeatN {
count: self.count,
element: self.element_ref().cloned().map_or_else(MaybeUninit::uninit, MaybeUninit::new),
}
}
}

#[stable(feature = "iter_repeat_n", since = "1.82.0")]
impl<A: fmt::Debug> fmt::Debug for RepeatN<A> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RepeatN")
.field("count", &self.count)
.field("element", &self.element_ref())
.finish()
}
}

#[stable(feature = "iter_repeat_n", since = "1.82.0")]
impl<A> Drop for RepeatN<A> {
fn drop(&mut self) {
Expand Down Expand Up @@ -194,9 +222,11 @@ impl<A: Clone> UncheckedIterator for RepeatN<A> {
// SAFETY: the check above ensured that the count used to be non-zero,
// so element hasn't been dropped yet, and we just lowered the count to
// zero so it won't be dropped later, and thus it's okay to take it here.
unsafe { ManuallyDrop::take(&mut self.element) }
unsafe { mem::replace(&mut self.element, MaybeUninit::uninit()).assume_init() }
} else {
A::clone(&self.element)
// SAFETY: the count is non-zero, so it must have not been dropped yet.
let element = unsafe { self.element.assume_init_ref() };
A::clone(element)
}
}
}
11 changes: 11 additions & 0 deletions library/core/tests/iter/sources.rs
Original file line number Diff line number Diff line change
Expand Up @@ -156,3 +156,14 @@ fn test_repeat_n_drop() {
drop((x0, x1, x2));
assert_eq!(count.get(), 3);
}

#[test]
fn test_repeat_n_soundness() {
let x = std::iter::repeat_n(String::from("use after free"), 0);
println!("{x:?}");

let mut y = std::iter::repeat_n(Box::new(0), 1);
let x = y.next().unwrap();
let _z = y;
assert_eq(0, *x);
}

0 comments on commit 361c3d3

Please sign in to comment.