@@ -64,27 +64,37 @@ where
6464 type Item = T ;
6565
6666 fn next ( & mut self ) -> Option < T > {
67- unsafe {
68- while self . idx < self . end {
69- let i = self . idx ;
70- let v = slice:: from_raw_parts_mut ( self . vec . as_mut_ptr ( ) , self . old_len ) ;
71- let drained = ( self . pred ) ( & mut v[ i] ) ;
72- // Update the index *after* the predicate is called. If the index
73- // is updated prior and the predicate panics, the element at this
74- // index would be leaked.
75- self . idx += 1 ;
76- if drained {
77- self . del += 1 ;
78- return Some ( ptr:: read ( & v[ i] ) ) ;
79- } else if self . del > 0 {
80- let del = self . del ;
81- let src: * const T = & v[ i] ;
82- let dst: * mut T = & mut v[ i - del] ;
83- ptr:: copy_nonoverlapping ( src, dst, 1 ) ;
67+ while self . idx < self . end {
68+ let i = self . idx ;
69+ // SAFETY:
70+ // We know that `i < self.end` from the if guard and that `self.end <= self.old_len` from
71+ // the validity of `Self`. Therefore `i` points to an element within `vec`.
72+ //
73+ // Additionally, the i-th element is valid because each element is visited at most once
74+ // and it is the first time we access vec[i].
75+ //
76+ // Note: we can't use `vec.get_unchecked_mut(i)` here since the precondition for that
77+ // function is that i < vec.len(), but we've set vec's length to zero.
78+ let cur = unsafe { & mut * self . vec . as_mut_ptr ( ) . add ( i) } ;
79+ let drained = ( self . pred ) ( cur) ;
80+ // Update the index *after* the predicate is called. If the index
81+ // is updated prior and the predicate panics, the element at this
82+ // index would be leaked.
83+ self . idx += 1 ;
84+ if drained {
85+ self . del += 1 ;
86+ // SAFETY: We never touch this element again after returning it.
87+ return Some ( unsafe { ptr:: read ( cur) } ) ;
88+ } else if self . del > 0 {
89+ // SAFETY: `self.del` > 0, so the hole slot must not overlap with current element.
90+ // We use copy for move, and never touch this element again.
91+ unsafe {
92+ let hole_slot = self . vec . as_mut_ptr ( ) . add ( i - self . del ) ;
93+ ptr:: copy_nonoverlapping ( cur, hole_slot, 1 ) ;
8494 }
8595 }
86- None
8796 }
97+ None
8898 }
8999
90100 fn size_hint ( & self ) -> ( usize , Option < usize > ) {
@@ -95,14 +105,18 @@ where
95105#[ stable( feature = "extract_if" , since = "1.87.0" ) ]
96106impl < T , F , A : Allocator > Drop for ExtractIf < ' _ , T , F , A > {
97107 fn drop ( & mut self ) {
98- unsafe {
99- if self . idx < self . old_len && self . del > 0 {
100- let ptr = self . vec . as_mut_ptr ( ) ;
101- let src = ptr. add ( self . idx ) ;
102- let dst = src. sub ( self . del ) ;
103- let tail_len = self . old_len - self . idx ;
104- src. copy_to ( dst, tail_len) ;
108+ if self . del > 0 {
109+ // SAFETY: Trailing unchecked items must be valid since we never touch them.
110+ unsafe {
111+ ptr:: copy (
112+ self . vec . as_ptr ( ) . add ( self . idx ) ,
113+ self . vec . as_mut_ptr ( ) . add ( self . idx - self . del ) ,
114+ self . old_len - self . idx ,
115+ ) ;
105116 }
117+ }
118+ // SAFETY: After filling holes, all items are in contiguous memory.
119+ unsafe {
106120 self . vec . set_len ( self . old_len - self . del ) ;
107121 }
108122 }
0 commit comments