diff --git a/library/alloc/src/vec/into_iter.rs b/library/alloc/src/vec/into_iter.rs index eae9ad076dca0..8a2d254a83451 100644 --- a/library/alloc/src/vec/into_iter.rs +++ b/library/alloc/src/vec/into_iter.rs @@ -164,18 +164,19 @@ impl Iterator for IntoIter { #[inline] fn advance_by(&mut self, n: usize) -> Result<(), usize> { let step_size = self.len().min(n); + let to_drop = ptr::slice_from_raw_parts_mut(self.ptr as *mut T, step_size); if mem::size_of::() == 0 { // SAFETY: due to unchecked casts of unsigned amounts to signed offsets the wraparound // effectively results in unsigned pointers representing positions 0..usize::MAX, // which is valid for ZSTs. self.ptr = unsafe { arith_offset(self.ptr as *const i8, step_size as isize) as *mut T } } else { - let to_drop = ptr::slice_from_raw_parts_mut(self.ptr as *mut T, step_size); // SAFETY: the min() above ensures that step_size is in bounds - unsafe { - self.ptr = self.ptr.add(step_size); - ptr::drop_in_place(to_drop); - } + self.ptr = unsafe { self.ptr.add(step_size) }; + } + // SAFETY: the min() above ensures that step_size is in bounds + unsafe { + ptr::drop_in_place(to_drop); } if step_size < n { return Err(step_size); @@ -237,11 +238,11 @@ impl DoubleEndedIterator for IntoIter { } else { // SAFETY: same as for advance_by() self.end = unsafe { self.end.offset(step_size.wrapping_neg() as isize) }; - let to_drop = ptr::slice_from_raw_parts_mut(self.end as *mut T, step_size); - // SAFETY: same as for advance_by() - unsafe { - ptr::drop_in_place(to_drop); - } + } + let to_drop = ptr::slice_from_raw_parts_mut(self.end as *mut T, step_size); + // SAFETY: same as for advance_by() + unsafe { + ptr::drop_in_place(to_drop); } if step_size < n { return Err(step_size); diff --git a/library/core/src/iter/adapters/flatten.rs b/library/core/src/iter/adapters/flatten.rs index e1d665bb2a177..351fd569d8acf 100644 --- a/library/core/src/iter/adapters/flatten.rs +++ b/library/core/src/iter/adapters/flatten.rs @@ -412,8 +412,9 @@ where self.frontiter = None; if let Some(ref mut back) = self.backiter { - if let Err(advanced) = back.advance_by(rem) { - rem -= advanced + match back.advance_by(rem) { + ret @ Ok(_) => return ret, + Err(advanced) => rem -= advanced, } } diff --git a/library/core/src/iter/adapters/skip.rs b/library/core/src/iter/adapters/skip.rs index e29ff1291cf8d..9b89ca5a9479d 100644 --- a/library/core/src/iter/adapters/skip.rs +++ b/library/core/src/iter/adapters/skip.rs @@ -116,14 +116,35 @@ where } #[inline] + #[rustc_inherit_overflow_checks] fn advance_by(&mut self, n: usize) -> Result<(), usize> { - if self.n >= n { - self.n -= n; - return Ok(()); + let mut rem = n; + + let step_one = self.n.saturating_add(rem); + match self.iter.advance_by(step_one) { + Ok(_) => { + rem -= step_one - self.n; + self.n = 0; + } + Err(advanced) => { + let advanced_without_skip = advanced.saturating_sub(self.n); + self.n = self.n.saturating_sub(advanced); + return Err(advanced_without_skip); + } } - let rem = n - self.n; - self.n = 0; - self.iter.advance_by(rem) + + // step_one calculation may have saturated + if unlikely(rem > 0) { + return match self.iter.advance_by(rem) { + ret @ Ok(_) => ret, + Err(advanced) => { + rem -= advanced; + Err(n - rem) + } + }; + } + + Ok(()) } }