1+ use crate :: fmt;
12use crate :: iter:: { FusedIterator , TrustedLen , UncheckedIterator } ;
2- use crate :: mem:: ManuallyDrop ;
3+ use crate :: mem:: { self , MaybeUninit } ;
34use crate :: num:: NonZero ;
45
56/// Creates a new iterator that repeats a single element a given number of times.
@@ -58,14 +59,12 @@ use crate::num::NonZero;
5859#[ inline]
5960#[ stable( feature = "iter_repeat_n" , since = "1.82.0" ) ]
6061pub fn repeat_n < T : Clone > ( element : T , count : usize ) -> RepeatN < T > {
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+ } ;
6968
7069 RepeatN { element, count }
7170}
@@ -74,31 +73,60 @@ pub fn repeat_n<T: Clone>(element: T, count: usize) -> RepeatN<T> {
7473///
7574/// This `struct` is created by the [`repeat_n()`] function.
7675/// See its documentation for more.
77- #[ derive( Clone , Debug ) ]
7876#[ stable( feature = "iter_repeat_n" , since = "1.82.0" ) ]
7977pub struct RepeatN < A > {
8078 count : usize ,
81- // Invariant: has been dropped iff count == 0.
82- element : ManuallyDrop < A > ,
79+ // Invariant: uninit iff count == 0.
80+ element : MaybeUninit < A > ,
8381}
8482
8583impl < A > RepeatN < A > {
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+ }
8693 /// If we haven't already dropped the element, return it in an option.
8794///
8895/// Clears the count so it won't be dropped again later.
8996#[ inline]
9097 fn take_element ( & mut self ) -> Option < A > {
9198 if self . count > 0 {
9299 self . count = 0 ;
100+ let element = mem:: replace ( & mut self . element , MaybeUninit :: uninit ( ) ) ;
93101 // SAFETY: We just set count to zero so it won't be dropped again,
94102 // 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 ( ) ) }
96104 } else {
97105 None
98106 }
99107 }
100108}
101109
110+ #[ stable( feature = "iter_repeat_n" , since = "1.82.0" ) ]
111+ impl < A : Clone > Clone for RepeatN < A > {
112+ fn clone ( & self ) -> RepeatN < A > {
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 < A > {
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+
102130#[ stable( feature = "iter_repeat_n" , since = "1.82.0" ) ]
103131impl < A > Drop for RepeatN < A > {
104132 fn drop ( & mut self ) {
@@ -194,9 +222,11 @@ impl<A: Clone> UncheckedIterator for RepeatN<A> {
194222 // SAFETY: the check above ensured that the count used to be non-zero,
195223 // so element hasn't been dropped yet, and we just lowered the count to
196224 // 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 ( ) }
198226 } 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)
200230 }
201231 }
202232}
0 commit comments