Auto merge of #130145 - fee1-dead-contrib:repeatn, r=lcnr,workingjubilee · qinheping/verify-rust-std@8a5922f (original) (raw)
``
1
`+
use crate::fmt;
`
1
2
`use crate::iter::{FusedIterator, TrustedLen, UncheckedIterator};
`
2
``
`-
use crate::mem::ManuallyDrop;
`
``
3
`+
use crate::mem::{self, MaybeUninit};
`
3
4
`use crate::num::NonZero;
`
4
5
``
5
6
`/// Creates a new iterator that repeats a single element a given number of times.
`
`@@ -58,14 +59,12 @@ use crate::num::NonZero;
`
58
59
`#[inline]
`
59
60
`#[stable(feature = "iter_repeat_n", since = "1.82.0")]
`
60
61
`pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN {
`
61
``
`-
let mut element = ManuallyDrop::new(element);
`
62
``
-
63
``
`-
if count == 0 {
`
64
``
`-
// SAFETY: we definitely haven't dropped it yet, since we only just got
`
65
``
`-
// passed it in, and because the count is zero the instance we're about
`
66
``
`-
// to create won't drop it, so to avoid leaking we need to now.
`
67
``
`-
unsafe { ManuallyDrop::drop(&mut element) };
`
68
``
`-
}
`
``
62
`+
let element = if count == 0 {
`
``
63
`` +
// element gets dropped eagerly.
``
``
64
`+
MaybeUninit::uninit()
`
``
65
`+
} else {
`
``
66
`+
MaybeUninit::new(element)
`
``
67
`+
};
`
69
68
``
70
69
`RepeatN { element, count }
`
71
70
`}
`
`@@ -74,31 +73,60 @@ pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN {
`
74
73
`///
`
75
74
`` /// This struct is created by the [repeat_n()] function.
``
76
75
`/// See its documentation for more.
`
77
``
`-
#[derive(Clone, Debug)]
`
78
76
`#[stable(feature = "iter_repeat_n", since = "1.82.0")]
`
79
77
`pub struct RepeatN {
`
80
78
`count: usize,
`
81
``
`-
// Invariant: has been dropped iff count == 0.
`
82
``
`-
element: ManuallyDrop,
`
``
79
`+
// Invariant: uninit iff count == 0.
`
``
80
`+
element: MaybeUninit,
`
83
81
`}
`
84
82
``
85
83
`
``
84
`+
/// Returns the element if it hasn't been dropped already.
`
``
85
`+
fn element_ref(&self) -> Option<&A> {
`
``
86
`+
if self.count > 0 {
`
``
87
`+
// SAFETY: The count is non-zero, so it must be initialized.
`
``
88
`+
Some(unsafe { self.element.assume_init_ref() })
`
``
89
`+
} else {
`
``
90
`+
None
`
``
91
`+
}
`
``
92
`+
}
`
86
93
`/// If we haven't already dropped the element, return it in an option.
`
87
94
`///
`
88
95
`/// Clears the count so it won't be dropped again later.
`
89
96
`#[inline]
`
90
97
`fn take_element(&mut self) -> Option {
`
91
98
`if self.count > 0 {
`
92
99
`self.count = 0;
`
``
100
`+
let element = mem::replace(&mut self.element, MaybeUninit::uninit());
`
93
101
`// SAFETY: We just set count to zero so it won't be dropped again,
`
94
102
`// and it used to be non-zero so it hasn't already been dropped.
`
95
``
`-
unsafe { Some(ManuallyDrop::take(&mut self.element)) }
`
``
103
`+
unsafe { Some(element.assume_init()) }
`
96
104
`} else {
`
97
105
`None
`
98
106
`}
`
99
107
`}
`
100
108
`}
`
101
109
``
``
110
`+
#[stable(feature = "iter_repeat_n", since = "1.82.0")]
`
``
111
`+
impl<A: Clone> Clone for RepeatN {
`
``
112
`+
fn clone(&self) -> RepeatN {
`
``
113
`+
RepeatN {
`
``
114
`+
count: self.count,
`
``
115
`+
element: self.element_ref().cloned().map_or_else(MaybeUninit::uninit, MaybeUninit::new),
`
``
116
`+
}
`
``
117
`+
}
`
``
118
`+
}
`
``
119
+
``
120
`+
#[stable(feature = "iter_repeat_n", since = "1.82.0")]
`
``
121
`+
impl<A: fmt::Debug> fmt::Debug for RepeatN {
`
``
122
`+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
`
``
123
`+
f.debug_struct("RepeatN")
`
``
124
`+
.field("count", &self.count)
`
``
125
`+
.field("element", &self.element_ref())
`
``
126
`+
.finish()
`
``
127
`+
}
`
``
128
`+
}
`
``
129
+
102
130
`#[stable(feature = "iter_repeat_n", since = "1.82.0")]
`
103
131
`impl Drop for RepeatN {
`
104
132
`fn drop(&mut self) {
`
`@@ -194,9 +222,11 @@ impl<A: Clone> UncheckedIterator for RepeatN {
`
194
222
`// SAFETY: the check above ensured that the count used to be non-zero,
`
195
223
`// so element hasn't been dropped yet, and we just lowered the count to
`
196
224
`// zero so it won't be dropped later, and thus it's okay to take it here.
`
197
``
`-
unsafe { ManuallyDrop::take(&mut self.element) }
`
``
225
`+
unsafe { mem::replace(&mut self.element, MaybeUninit::uninit()).assume_init() }
`
198
226
`} else {
`
199
``
`-
A::clone(&self.element)
`
``
227
`+
// SAFETY: the count is non-zero, so it must have not been dropped yet.
`
``
228
`+
let element = unsafe { self.element.assume_init_ref() };
`
``
229
`+
A::clone(element)
`
200
230
`}
`
201
231
`}
`
202
232
`}
`