Rollup merge of #138562 - kornelski:nth-panic, r=Noratrieb

Optimize slice {Chunks,Windows}::nth

I've noticed that the `nth` functions on slice iters had non-optimized-out bounds checks.

The new implementation even generates branchless code.
This commit is contained in:
Guillaume Gomez 2025-04-06 18:08:08 +02:00 committed by GitHub
commit 7bd89b90e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -1380,14 +1380,16 @@ impl<'a, T> Iterator for Windows<'a, T> {
#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let (end, overflow) = self.size.get().overflowing_add(n);
if end > self.v.len() || overflow {
self.v = &[];
None
} else {
let nth = &self.v[n..end];
self.v = &self.v[n + 1..];
let size = self.size.get();
if let Some(rest) = self.v.get(n..)
&& let Some(nth) = rest.get(..size)
{
self.v = &rest[1..];
Some(nth)
} else {
// setting length to 0 is cheaper than overwriting the pointer when assigning &[]
self.v = &self.v[..0]; // cheaper than &[]
None
}
}
@ -1427,7 +1429,7 @@ impl<'a, T> DoubleEndedIterator for Windows<'a, T> {
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let (end, overflow) = self.v.len().overflowing_sub(n);
if end < self.size.get() || overflow {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
let ret = &self.v[end - self.size.get()..end];
@ -1536,17 +1538,15 @@ impl<'a, T> Iterator for Chunks<'a, T> {
#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let (start, overflow) = n.overflowing_mul(self.chunk_size);
if start >= self.v.len() || overflow {
self.v = &[];
None
} else {
let end = match start.checked_add(self.chunk_size) {
Some(sum) => cmp::min(self.v.len(), sum),
None => self.v.len(),
};
let nth = &self.v[start..end];
self.v = &self.v[end..];
// min(len) makes a wrong start harmless, but enables optimizing this to brachless code
let chunk_start = &self.v[start.min(self.v.len())..];
let (nth, remainder) = chunk_start.split_at(self.chunk_size.min(chunk_start.len()));
if !overflow && start < self.v.len() {
self.v = remainder;
Some(nth)
} else {
self.v = &self.v[..0]; // cheaper than &[]
None
}
}
@ -1609,7 +1609,7 @@ impl<'a, T> DoubleEndedIterator for Chunks<'a, T> {
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let len = self.len();
if n >= len {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
let start = (len - 1 - n) * self.chunk_size;
@ -1933,7 +1933,7 @@ impl<'a, T> Iterator for ChunksExact<'a, T> {
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let (start, overflow) = n.overflowing_mul(self.chunk_size);
if start >= self.v.len() || overflow {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
let (_, snd) = self.v.split_at(start);
@ -1971,7 +1971,7 @@ impl<'a, T> DoubleEndedIterator for ChunksExact<'a, T> {
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let len = self.len();
if n >= len {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
let start = (len - 1 - n) * self.chunk_size;
@ -2638,7 +2638,7 @@ impl<'a, T> Iterator for RChunks<'a, T> {
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let (end, overflow) = n.overflowing_mul(self.chunk_size);
if end >= self.v.len() || overflow {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
// Can't underflow because of the check above
@ -2695,7 +2695,7 @@ impl<'a, T> DoubleEndedIterator for RChunks<'a, T> {
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let len = self.len();
if n >= len {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
// can't underflow because `n < len`
@ -3023,7 +3023,7 @@ impl<'a, T> Iterator for RChunksExact<'a, T> {
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let (end, overflow) = n.overflowing_mul(self.chunk_size);
if end >= self.v.len() || overflow {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
let (fst, _) = self.v.split_at(self.v.len() - end);
@ -3062,7 +3062,7 @@ impl<'a, T> DoubleEndedIterator for RChunksExact<'a, T> {
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let len = self.len();
if n >= len {
self.v = &[];
self.v = &self.v[..0]; // cheaper than &[]
None
} else {
// now that we know that `n` corresponds to a chunk,