Skip to content

Commit ffd7ade

Browse files
committed
fix issues pointed out in review
1 parent 6654a0b commit ffd7ade

File tree

3 files changed

+41
-18
lines changed

3 files changed

+41
-18
lines changed

library/alloc/src/vec/into_iter.rs

+11-10
Original file line numberDiff line numberDiff line change
@@ -164,18 +164,19 @@ impl<T, A: Allocator> Iterator for IntoIter<T, A> {
164164
#[inline]
165165
fn advance_by(&mut self, n: usize) -> Result<(), usize> {
166166
let step_size = self.len().min(n);
167+
let to_drop = ptr::slice_from_raw_parts_mut(self.ptr as *mut T, step_size);
167168
if mem::size_of::<T>() == 0 {
168169
// SAFETY: due to unchecked casts of unsigned amounts to signed offsets the wraparound
169170
// effectively results in unsigned pointers representing positions 0..usize::MAX,
170171
// which is valid for ZSTs.
171172
self.ptr = unsafe { arith_offset(self.ptr as *const i8, step_size as isize) as *mut T }
172173
} else {
173-
let to_drop = ptr::slice_from_raw_parts_mut(self.ptr as *mut T, step_size);
174174
// SAFETY: the min() above ensures that step_size is in bounds
175-
unsafe {
176-
self.ptr = self.ptr.add(step_size);
177-
ptr::drop_in_place(to_drop);
178-
}
175+
self.ptr = unsafe { self.ptr.add(step_size) };
176+
}
177+
// SAFETY: the min() above ensures that step_size is in bounds
178+
unsafe {
179+
ptr::drop_in_place(to_drop);
179180
}
180181
if step_size < n {
181182
return Err(step_size);
@@ -237,11 +238,11 @@ impl<T, A: Allocator> DoubleEndedIterator for IntoIter<T, A> {
237238
} else {
238239
// SAFETY: same as for advance_by()
239240
self.end = unsafe { self.end.offset(step_size.wrapping_neg() as isize) };
240-
let to_drop = ptr::slice_from_raw_parts_mut(self.end as *mut T, step_size);
241-
// SAFETY: same as for advance_by()
242-
unsafe {
243-
ptr::drop_in_place(to_drop);
244-
}
241+
}
242+
let to_drop = ptr::slice_from_raw_parts_mut(self.end as *mut T, step_size);
243+
// SAFETY: same as for advance_by()
244+
unsafe {
245+
ptr::drop_in_place(to_drop);
245246
}
246247
if step_size < n {
247248
return Err(step_size);

library/core/src/iter/adapters/flatten.rs

+3-2
Original file line numberDiff line numberDiff line change
@@ -412,8 +412,9 @@ where
412412
self.frontiter = None;
413413

414414
if let Some(ref mut back) = self.backiter {
415-
if let Err(advanced) = back.advance_by(rem) {
416-
rem -= advanced
415+
match back.advance_by(rem) {
416+
ret @ Ok(_) => return ret,
417+
Err(advanced) => rem -= advanced,
417418
}
418419
}
419420

library/core/src/iter/adapters/skip.rs

+27-6
Original file line numberDiff line numberDiff line change
@@ -116,14 +116,35 @@ where
116116
}
117117

118118
#[inline]
119+
#[rustc_inherit_overflow_checks]
119120
fn advance_by(&mut self, n: usize) -> Result<(), usize> {
120-
if self.n >= n {
121-
self.n -= n;
122-
return Ok(());
121+
let mut rem = n;
122+
123+
let step_one = self.n.saturating_add(rem);
124+
match self.iter.advance_by(step_one) {
125+
Ok(_) => {
126+
rem -= step_one - self.n;
127+
self.n = 0;
128+
}
129+
Err(advanced) => {
130+
let advanced_without_skip = advanced.saturating_sub(self.n);
131+
self.n = self.n.saturating_sub(advanced);
132+
return Err(advanced_without_skip);
133+
}
123134
}
124-
let rem = n - self.n;
125-
self.n = 0;
126-
self.iter.advance_by(rem)
135+
136+
// step_one calculation may have saturated
137+
if unlikely(rem > 0) {
138+
return match self.iter.advance_by(rem) {
139+
ret @ Ok(_) => ret,
140+
Err(advanced) => {
141+
rem -= advanced;
142+
Err(n - rem)
143+
}
144+
};
145+
}
146+
147+
Ok(())
127148
}
128149
}
129150

0 commit comments

Comments
 (0)