Skip to content

Commit 481bccb

Browse files
authored
Rollup merge of rust-lang#81617 - sdroege:vec-drain-trusted-random-access, r=Amanieu
Implement `TrustedRandomAccess` for `vec::Drain`
2 parents e11a9fa + 1d33db3 commit 481bccb

File tree

1 file changed

+45
-9
lines changed

1 file changed

+45
-9
lines changed

library/alloc/src/vec/drain.rs

Lines changed: 45 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
use crate::alloc::{Allocator, Global};
22
use core::fmt;
3-
use core::iter::{FusedIterator, TrustedLen};
3+
use core::iter::{FusedIterator, TrustedLen, TrustedRandomAccess};
44
use core::mem::{self};
55
use core::ptr::{self, NonNull};
66
use core::slice::{self};
@@ -89,6 +89,19 @@ impl<T, A: Allocator> Iterator for Drain<'_, T, A> {
8989
fn size_hint(&self) -> (usize, Option<usize>) {
9090
self.iter.size_hint()
9191
}
92+
93+
#[doc(hidden)]
94+
unsafe fn __iterator_get_unchecked(&mut self, idx: usize) -> T
95+
where
96+
Self: TrustedRandomAccess,
97+
{
98+
// SAFETY: `TrustedRandomAccess` requires that `idx` is in bounds and that
99+
// each `idx` is only accessed once. Forwarding to the slice iterator's
100+
// implementation is thus safe, and reading the value is safe because
101+
// `Self: TrustedRandomAccess` implies `T: Copy` so the `Drop` impl below
102+
// won't cause each item to be dropped twice.
103+
unsafe { ptr::read(self.iter.__iterator_get_unchecked(idx) as *const _) }
104+
}
92105
}
93106

94107
#[stable(feature = "drain", since = "1.6.0")]
@@ -108,9 +121,11 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
108121

109122
impl<'r, 'a, T, A: Allocator> Drop for DropGuard<'r, 'a, T, A> {
110123
fn drop(&mut self) {
111-
// Continue the same loop we have below. If the loop already finished, this does
112-
// nothing.
113-
self.0.for_each(drop);
124+
if mem::needs_drop::<T>() {
125+
// Continue the same loop we have below. If the loop already finished, this does
126+
// nothing.
127+
self.0.for_each(drop);
128+
}
114129

115130
if self.0.tail_len > 0 {
116131
unsafe {
@@ -129,11 +144,13 @@ impl<T, A: Allocator> Drop for Drain<'_, T, A> {
129144
}
130145
}
131146

132-
// exhaust self first
133-
while let Some(item) = self.next() {
134-
let guard = DropGuard(self);
135-
drop(item);
136-
mem::forget(guard);
147+
// exhaust self first if dropping of the items is required
148+
if mem::needs_drop::<T>() {
149+
while let Some(item) = self.next() {
150+
let guard = DropGuard(self);
151+
drop(item);
152+
mem::forget(guard);
153+
}
137154
}
138155

139156
// Drop a `DropGuard` to move back the non-drained tail of `self`.
@@ -149,7 +166,26 @@ impl<T, A: Allocator> ExactSizeIterator for Drain<'_, T, A> {
149166
}
150167

151168
#[unstable(feature = "trusted_len", issue = "37572")]
169+
// SAFETY: `Drain` simply forwards to the underlying slice iterator, which implements `TrustedLen`
170+
// so the required properties are all preserved.
152171
unsafe impl<T, A: Allocator> TrustedLen for Drain<'_, T, A> {}
153172

173+
#[doc(hidden)]
174+
#[unstable(feature = "trusted_random_access", issue = "none")]
175+
// SAFETY: `Drain` forwards to the underlying slice iterator, which implements `TrustedRandomAccess`,
176+
// and then reads the items instead of just returning a reference. As `TrustedRandomAccess`
177+
// requires each index to be accessed only once, this is safe to do here.
178+
//
179+
// T: Copy as approximation for !Drop since get_unchecked does not advance self.iter
180+
// and as a result the `Drop` impl above would otherwise cause items to be dropped twice.
181+
unsafe impl<T, A: Allocator> TrustedRandomAccess for Drain<'_, T, A>
182+
where
183+
T: Copy,
184+
{
185+
fn may_have_side_effect() -> bool {
186+
false
187+
}
188+
}
189+
154190
#[stable(feature = "fused", since = "1.26.0")]
155191
impl<T, A: Allocator> FusedIterator for Drain<'_, T, A> {}

0 commit comments

Comments
 (0)