Skip to content

Commit 2da2fab

Browse files
committed
Add slice::DrainRaw that moves out of a NonNull<[T]>
1 parent a9034ef commit 2da2fab

File tree

5 files changed

+287
-18
lines changed

5 files changed

+287
-18
lines changed

library/core/src/lib.rs

+1
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,7 @@
224224
#![feature(doc_cfg)]
225225
#![feature(doc_cfg_hide)]
226226
#![feature(doc_notable_trait)]
227+
#![feature(dropck_eyepatch)]
227228
#![feature(effects)]
228229
#![feature(extern_types)]
229230
#![feature(f128)]

library/core/src/slice/drain.rs

+250
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
use crate::array;
2+
use crate::fmt;
3+
use crate::iter::{
4+
FusedIterator, TrustedFused, TrustedLen, TrustedRandomAccessNoCoerce, UncheckedIterator,
5+
};
6+
use crate::mem::MaybeUninit;
7+
use crate::num::NonZero;
8+
use crate::ptr::NonNull;
9+
use crate::slice::NonNullIter;
10+
11+
/// An iterator which takes ownership of items out of a slice, dropping any
12+
/// remaining items when the iterator drops.
13+
///
14+
/// Note that, like a raw pointer, it's **up to you** to get the lifetime right.
15+
/// In some ways it's actually harder to get right, as the iterator interface
16+
/// appears safe, but as you promise when creating one of these, you still must
17+
/// ensure that the mentioned memory is usable the whole time this lives.
18+
///
19+
/// Ideally you won't be using this directly, but rather a version encapsulated
20+
/// in a safer interface, like `vec::IntoIter`.
21+
///
22+
/// This raw version may be removed in favour of a future language feature,
23+
/// such as using `unsafe<'a> Drain<'a, T>` instead of `DrainRaw<T>`.
24+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
25+
pub struct DrainRaw<T>(NonNullIter<T>);
26+
27+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
28+
// `may_dangle` is needed for compatibility with `vec::IntoIter`
29+
unsafe impl<#[may_dangle] T> Drop for DrainRaw<T> {
30+
fn drop(&mut self) {
31+
// When used in things like `vec::IntoIter`, the memory over which we're
32+
// iterating might have been deallocated once we're running this drop.
33+
// At the time of writing, Miri doesn't like `sub_ptr` between pointers
34+
// into a deallocated allocation. So checking empty first -- which just
35+
// needs pointer equality -- avoids that issue.
36+
if !self.is_empty() {
37+
let slice = self.as_nonnull_slice();
38+
// SAFETY: By type invariant, we're allowed to drop the rest of the items.
39+
unsafe { slice.drop_in_place() };
40+
}
41+
}
42+
}
43+
44+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
45+
impl<T: fmt::Debug> fmt::Debug for DrainRaw<T> {
46+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47+
f.debug_tuple("DrainRaw").field(&self.0.make_shortlived_slice()).finish()
48+
}
49+
}
50+
51+
impl<T> DrainRaw<T> {
52+
/// Creates a new iterator which moves the `len` items starting at `ptr`
53+
/// while it's iterated, or drops them when the iterator is dropped.
54+
///
55+
/// # Safety
56+
///
57+
/// - `ptr` through `ptr.add(len)` must be a single allocated object
58+
/// such that that it's sound to `offset` through it.
59+
/// - All those elements must be readable, including being sufficiently aligned.
60+
/// - All those elements are valid for dropping.
61+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
62+
#[inline]
63+
pub unsafe fn from_parts(ptr: NonNull<T>, len: usize) -> Self {
64+
// SAFETY: this function's safety conditions are stricter than NonNullIter,
65+
// and include allowing the type to drop the items in `Drop`.
66+
Self(unsafe { NonNullIter::from_parts(ptr, len) })
67+
}
68+
69+
/// Returns a pointer to the remaining elements of the iterator
70+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
71+
#[inline]
72+
pub fn as_nonnull_slice(&self) -> NonNull<[T]> {
73+
self.0.make_nonnull_slice()
74+
}
75+
76+
/// Equivalent to exhausting the iterator normally, but faster.
77+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
78+
#[inline]
79+
pub fn drop_remaining(&mut self) {
80+
let all = self.forget_remaining();
81+
// SAFETY: We "forgot" these elements so our `Drop` won't drop them,
82+
// so it's ok to drop them here without risking double-frees.
83+
unsafe { all.drop_in_place() }
84+
}
85+
86+
/// Exhaust the iterator without actually dropping the rest of the items.
87+
///
88+
/// Returns the forgotten items.
89+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
90+
#[inline]
91+
pub fn forget_remaining(&mut self) -> NonNull<[T]> {
92+
let all = self.as_nonnull_slice();
93+
self.0.exhaust();
94+
all
95+
}
96+
}
97+
98+
impl<T> UncheckedIterator for DrainRaw<T> {
99+
#[inline]
100+
unsafe fn next_unchecked(&mut self) -> T {
101+
// SAFETY: we're a 1:1 mapping of the inner iterator, so if the caller
102+
// proved we have another item, the inner iterator has another one too.
103+
// Also, the `next_unchecked` means the returned item is no longer part
104+
// of the inner iterator, and thus `read`ing it here -- and giving it
105+
// to the caller who will (probably) drop it -- is ok.
106+
unsafe { self.0.next_unchecked().read() }
107+
}
108+
}
109+
110+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
111+
impl<T> Iterator for DrainRaw<T> {
112+
type Item = T;
113+
114+
#[inline]
115+
fn next(&mut self) -> Option<T> {
116+
match self.0.next() {
117+
// SAFETY: The `next` means the returned item is no longer part of
118+
// the inner iterator, and thus `read`ing it here -- and giving it
119+
// to the caller who will (probably) drop it -- is ok.
120+
Some(ptr) => Some(unsafe { ptr.read() }),
121+
None => None,
122+
}
123+
}
124+
125+
#[inline]
126+
fn size_hint(&self) -> (usize, Option<usize>) {
127+
self.0.size_hint()
128+
}
129+
130+
#[inline]
131+
fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
132+
let clamped = self.len().min(n);
133+
// SAFETY: By construction, `clamped` is always in-bounds.
134+
// The skipped elements are removed from the inner iterator so won't be
135+
// dropped in `Drop`, so dropping there here is fine.
136+
unsafe {
137+
let to_drop = self.0.skip_forward_unchecked(clamped);
138+
to_drop.drop_in_place();
139+
}
140+
NonZero::new(n - clamped).map_or(Ok(()), Err)
141+
}
142+
143+
#[inline]
144+
fn count(self) -> usize {
145+
self.len()
146+
}
147+
148+
#[inline]
149+
fn next_chunk<const N: usize>(&mut self) -> Result<[T; N], core::array::IntoIter<T, N>> {
150+
let len = self.len();
151+
let clamped = len.min(N);
152+
153+
// SAFETY: By construction, `clamped` is always in-bounds.
154+
let to_copy = unsafe { self.0.skip_forward_unchecked(clamped) };
155+
if len >= N {
156+
// SAFETY: If we have more elements than were requested, they can be
157+
// read directly because arrays need no extra alignment.
158+
Ok(unsafe { to_copy.cast::<[T; N]>().read() })
159+
} else {
160+
let mut raw_ary = MaybeUninit::uninit_array();
161+
// SAFETY: If we don't have enough elements left, then copy all the
162+
// ones we do have into the local array, which cannot overlap because
163+
// new locals are always distinct storage.
164+
Err(unsafe {
165+
MaybeUninit::<T>::slice_as_mut_ptr(&mut raw_ary)
166+
.copy_from_nonoverlapping(to_copy.as_mut_ptr(), len);
167+
array::IntoIter::new_unchecked(raw_ary, 0..len)
168+
})
169+
}
170+
}
171+
172+
unsafe fn __iterator_get_unchecked(&mut self, i: usize) -> Self::Item
173+
where
174+
Self: TrustedRandomAccessNoCoerce,
175+
{
176+
// SAFETY: the caller must guarantee that `i` is in bounds of the slice,
177+
// so the `get_unchecked_mut(i)` is guaranteed to pointer to an element
178+
// and thus guaranteed to be valid to dereference.
179+
//
180+
// Also note the implementation of `Self: TrustedRandomAccess` requires
181+
// that `T: Copy` so reading elements from the buffer doesn't invalidate
182+
// them for `Drop`.
183+
unsafe { self.as_nonnull_slice().get_unchecked_mut(i).read() }
184+
}
185+
}
186+
187+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
188+
impl<T> DoubleEndedIterator for DrainRaw<T> {
189+
#[inline]
190+
fn next_back(&mut self) -> Option<T> {
191+
match self.0.next_back() {
192+
// SAFETY: The `next_back` means the returned item is no longer part of
193+
// the inner iterator, and thus `read`ing it here -- and giving it
194+
// to the caller who will (probably) drop it -- is ok.
195+
Some(ptr) => Some(unsafe { ptr.read() }),
196+
None => None,
197+
}
198+
}
199+
200+
#[inline]
201+
fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
202+
let clamped = self.len().min(n);
203+
// SAFETY: By construction, `clamped` is always in-bounds.
204+
// The skipped elements are removed from the inner iterator so won't be
205+
// dropped in `Drop`, so dropping there here is fine.
206+
unsafe {
207+
let to_drop = self.0.skip_backward_unchecked(clamped);
208+
to_drop.drop_in_place();
209+
}
210+
NonZero::new(n - clamped).map_or(Ok(()), Err)
211+
}
212+
}
213+
214+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
215+
impl<T> ExactSizeIterator for DrainRaw<T> {
216+
fn is_empty(&self) -> bool {
217+
self.0.is_empty()
218+
}
219+
220+
fn len(&self) -> usize {
221+
self.0.len()
222+
}
223+
}
224+
225+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
226+
impl<T> FusedIterator for DrainRaw<T> {}
227+
228+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
229+
#[doc(hidden)]
230+
unsafe impl<T> TrustedFused for DrainRaw<T> {}
231+
232+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
233+
unsafe impl<T> TrustedLen for DrainRaw<T> {}
234+
235+
#[doc(hidden)]
236+
#[unstable(issue = "none", feature = "std_internals")]
237+
#[rustc_unsafe_specialization_marker]
238+
pub trait NonDrop {}
239+
240+
// T: Copy as approximation for !Drop since get_unchecked does not advance self.ptr
241+
// and thus we can't implement drop-handling
242+
#[unstable(issue = "none", feature = "std_internals")]
243+
impl<T: Copy> NonDrop for T {}
244+
245+
// SAFETY: If they're accessing things in random order we don't know when to drop
246+
// things, so only allow this for `Copy` things where that doesn't matter.
247+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
248+
unsafe impl<T: NonDrop> TrustedRandomAccessNoCoerce for DrainRaw<T> {
249+
const MAY_HAVE_SIDE_EFFECT: bool = false;
250+
}

library/core/src/slice/iter.rs

+10-16
Original file line numberDiff line numberDiff line change
@@ -411,7 +411,6 @@ impl<T> AsRef<[T]> for IterMut<'_, T> {
411411
iterator! {struct IterMut<'a, T> => *mut T, &'a mut T, {}}
412412

413413
/// Iterator over all the `NonNull<T>` pointers to the elements of a slice.
414-
#[unstable(feature = "slice_non_null_iter", issue = "none")]
415414
#[must_use = "iterators are lazy and do nothing unless consumed"]
416415
pub struct NonNullIter<T> {
417416
/// The pointer to the next element to return, or the past-the-end location
@@ -425,26 +424,13 @@ pub struct NonNullIter<T> {
425424
end_or_len: *const T,
426425
}
427426

428-
#[unstable(feature = "slice_non_null_iter", issue = "none")]
429427
impl<T: fmt::Debug> fmt::Debug for NonNullIter<T> {
430428
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
431429
f.debug_tuple("NonNullIter").field(&self.make_shortlived_slice()).finish()
432430
}
433431
}
434432

435433
impl<T> NonNullIter<T> {
436-
/// Turn an iterator giving `&T`s into one giving `NonNull<T>`s.
437-
#[unstable(feature = "slice_non_null_iter", issue = "none")]
438-
pub fn from_slice_iter(Iter { ptr, end_or_len, .. }: Iter<'_, T>) -> Self {
439-
Self { ptr, end_or_len }
440-
}
441-
442-
/// Turn an iterator giving `&mut T`s into one giving `NonNull<T>`s.
443-
#[unstable(feature = "slice_non_null_iter", issue = "none")]
444-
pub fn from_slice_iter_mut(IterMut { ptr, end_or_len, .. }: IterMut<'_, T>) -> Self {
445-
Self { ptr, end_or_len }
446-
}
447-
448434
/// Creates a new iterator over the `len` items starting at `ptr`
449435
///
450436
/// # Safety
@@ -453,7 +439,6 @@ impl<T> NonNullIter<T> {
453439
/// such that that it's sound to `offset` through it.
454440
/// - All those elements must be readable
455441
/// - The caller must ensure both as long as the iterator is in use.
456-
#[unstable(feature = "slice_non_null_iter", issue = "none")]
457442
#[inline]
458443
pub unsafe fn from_parts(ptr: NonNull<T>, len: usize) -> Self {
459444
// SAFETY: There are several things here:
@@ -481,7 +466,16 @@ impl<T> NonNullIter<T> {
481466
}
482467

483468
#[inline]
484-
unsafe fn non_null_to_item(p: NonNull<T>) -> <Self as Iterator>::Item {
469+
pub fn exhaust(&mut self) {
470+
if T::IS_ZST {
471+
self.end_or_len = without_provenance_mut(0);
472+
} else {
473+
self.end_or_len = self.ptr.as_ptr();
474+
}
475+
}
476+
477+
#[inline]
478+
fn non_null_to_item(p: NonNull<T>) -> <Self as Iterator>::Item {
485479
p
486480
}
487481
}

library/core/src/slice/iter/macros.rs

+20
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,26 @@ macro_rules! iterator {
134134
},
135135
)
136136
}
137+
138+
// This is not used on every type that uses this macro, but is more
139+
// convenient to implement here so it can use `post_inc_start`.
140+
#[allow(dead_code)]
141+
#[inline]
142+
pub(crate) unsafe fn skip_forward_unchecked(&mut self, offset: usize) -> NonNull<[T]> {
143+
// SAFETY: The caller guarantees the provided offset is in-bounds.
144+
let old_begin = unsafe { self.post_inc_start(offset) };
145+
NonNull::slice_from_raw_parts(old_begin, offset)
146+
}
147+
148+
// This is not used on every type that uses this macro, but is more
149+
// convenient to implement here so it can use `pre_dec_end`.
150+
#[allow(dead_code)]
151+
#[inline]
152+
pub(crate) unsafe fn skip_backward_unchecked(&mut self, offset: usize) -> NonNull<[T]> {
153+
// SAFETY: The caller guarantees the provided offset is in-bounds.
154+
let new_end = unsafe { self.pre_dec_end(offset) };
155+
NonNull::slice_from_raw_parts(new_end, offset)
156+
}
137157
}
138158

139159
#[allow(unused_lifetimes)]

library/core/src/slice/mod.rs

+6-2
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ pub mod sort;
3535

3636
mod ascii;
3737
mod cmp;
38+
mod drain;
3839
pub(crate) mod index;
3940
mod iter;
4041
mod raw;
@@ -46,8 +47,6 @@ mod specialize;
4647
#[doc(hidden)]
4748
pub use ascii::is_ascii_simple;
4849

49-
#[unstable(feature = "slice_non_null_iter", issue = "none")]
50-
pub use iter::NonNullIter;
5150
#[stable(feature = "rust1", since = "1.0.0")]
5251
pub use iter::{Chunks, ChunksMut, Windows};
5352
#[stable(feature = "rust1", since = "1.0.0")]
@@ -73,6 +72,11 @@ pub use iter::ArrayWindows;
7372
#[stable(feature = "slice_group_by", since = "1.77.0")]
7473
pub use iter::{ChunkBy, ChunkByMut};
7574

75+
use iter::NonNullIter;
76+
77+
#[unstable(feature = "slice_drain_raw_iter", issue = "none")]
78+
pub use drain::DrainRaw;
79+
7680
#[stable(feature = "split_inclusive", since = "1.51.0")]
7781
pub use iter::{SplitInclusive, SplitInclusiveMut};
7882

0 commit comments

Comments
 (0)