Skip to content

Commit 3fcbc31

Browse files
committed
Optimize iterator adapters.
Specifically, make count, nth, and last call the corresponding methods on the underlying iterator where possible. This way, if the underlying iterator has an optimized count, nth, or last implementations (e.g. slice::Iter), these methods will propagate these optimizations. Additionally, change Skip::next to take advantage of a potentially optimized nth method on the underlying iterator.
1 parent 6cd7486 commit 3fcbc31

File tree

2 files changed

+360
-18
lines changed

2 files changed

+360
-18
lines changed

src/libcore/iter.rs

+145-18
Original file line numberDiff line numberDiff line change
@@ -1472,6 +1472,32 @@ impl<A, B> Iterator for Chain<A, B> where
14721472
}
14731473
}
14741474

1475+
#[inline]
1476+
fn count(self) -> usize {
1477+
(if !self.flag { self.a.count() } else { 0 }) + self.b.count()
1478+
}
1479+
1480+
#[inline]
1481+
fn nth(&mut self, mut n: usize) -> Option<A::Item> {
1482+
if !self.flag {
1483+
for x in self.a.by_ref() {
1484+
if n == 0 {
1485+
return Some(x)
1486+
}
1487+
n -= 1;
1488+
}
1489+
self.flag = true;
1490+
}
1491+
self.b.nth(n)
1492+
}
1493+
1494+
#[inline]
1495+
fn last(self) -> Option<A::Item> {
1496+
let a_last = if self.flag { None } else { self.a.last() };
1497+
let b_last = self.b.last();
1498+
b_last.or(a_last)
1499+
}
1500+
14751501
#[inline]
14761502
fn size_hint(&self) -> (usize, Option<usize>) {
14771503
let (a_lower, a_upper) = self.a.size_hint();
@@ -1777,6 +1803,20 @@ impl<I> Iterator for Enumerate<I> where I: Iterator {
17771803
fn size_hint(&self) -> (usize, Option<usize>) {
17781804
self.iter.size_hint()
17791805
}
1806+
1807+
#[inline]
1808+
fn nth(&mut self, n: usize) -> Option<(usize, I::Item)> {
1809+
self.iter.nth(n).map(|a| {
1810+
let i = self.count + n;
1811+
self.count = i + 1;
1812+
(i, a)
1813+
})
1814+
}
1815+
1816+
#[inline]
1817+
fn count(self) -> usize {
1818+
self.iter.count()
1819+
}
17801820
}
17811821

17821822
#[stable(feature = "rust1", since = "1.0.0")]
@@ -1834,6 +1874,28 @@ impl<I: Iterator> Iterator for Peekable<I> {
18341874
}
18351875
}
18361876

1877+
#[inline]
1878+
fn count(self) -> usize {
1879+
(if self.peeked.is_some() { 1 } else { 0 }) + self.iter.count()
1880+
}
1881+
1882+
#[inline]
1883+
fn nth(&mut self, n: usize) -> Option<I::Item> {
1884+
match self.peeked {
1885+
Some(_) if n == 0 => self.peeked.take(),
1886+
Some(_) => {
1887+
self.peeked = None;
1888+
self.iter.nth(n-1)
1889+
},
1890+
None => self.iter.nth(n)
1891+
}
1892+
}
1893+
1894+
#[inline]
1895+
fn last(self) -> Option<I::Item> {
1896+
self.iter.last().or(self.peeked)
1897+
}
1898+
18371899
#[inline]
18381900
fn size_hint(&self) -> (usize, Option<usize>) {
18391901
let (lo, hi) = self.iter.size_hint();
@@ -1960,27 +2022,49 @@ impl<I> Iterator for Skip<I> where I: Iterator {
19602022
type Item = <I as Iterator>::Item;
19612023

19622024
#[inline]
1963-
fn next(&mut self) -> Option<<I as Iterator>::Item> {
1964-
let mut next = self.iter.next();
2025+
fn next(&mut self) -> Option<I::Item> {
19652026
if self.n == 0 {
1966-
next
2027+
self.iter.next()
19672028
} else {
1968-
let mut n = self.n;
1969-
while n > 0 {
1970-
n -= 1;
1971-
match next {
1972-
Some(_) => {
1973-
next = self.iter.next();
1974-
continue
1975-
}
1976-
None => {
1977-
self.n = 0;
1978-
return None
1979-
}
1980-
}
1981-
}
2029+
let old_n = self.n;
19822030
self.n = 0;
1983-
next
2031+
self.iter.nth(old_n)
2032+
}
2033+
}
2034+
2035+
#[inline]
2036+
fn nth(&mut self, n: usize) -> Option<I::Item> {
2037+
// Can't just add n + self.n due to overflow.
2038+
if self.n == 0 {
2039+
self.iter.nth(n)
2040+
} else {
2041+
let to_skip = self.n;
2042+
self.n = 0;
2043+
// nth(n) skips n+1
2044+
if self.iter.nth(to_skip-1).is_none() {
2045+
return None;
2046+
}
2047+
self.iter.nth(n)
2048+
}
2049+
}
2050+
2051+
#[inline]
2052+
fn count(self) -> usize {
2053+
self.iter.count().saturating_sub(self.n)
2054+
}
2055+
2056+
#[inline]
2057+
fn last(mut self) -> Option<I::Item> {
2058+
if self.n == 0 {
2059+
self.iter.last()
2060+
} else {
2061+
let next = self.next();
2062+
if next.is_some() {
2063+
// recurse. n should be 0.
2064+
self.last().or(next)
2065+
} else {
2066+
None
2067+
}
19842068
}
19852069
}
19862070

@@ -2038,6 +2122,20 @@ impl<I> Iterator for Take<I> where I: Iterator{
20382122
}
20392123
}
20402124

2125+
#[inline]
2126+
fn nth(&mut self, n: usize) -> Option<I::Item> {
2127+
if self.n > n {
2128+
self.n -= n + 1;
2129+
self.iter.nth(n)
2130+
} else {
2131+
if self.n > 0 {
2132+
self.iter.nth(self.n - 1);
2133+
self.n = 0;
2134+
}
2135+
None
2136+
}
2137+
}
2138+
20412139
#[inline]
20422140
fn size_hint(&self) -> (usize, Option<usize>) {
20432141
let (lower, upper) = self.iter.size_hint();
@@ -2199,6 +2297,35 @@ impl<I> Iterator for Fuse<I> where I: Iterator {
21992297
}
22002298
}
22012299

2300+
#[inline]
2301+
fn nth(&mut self, n: usize) -> Option<I::Item> {
2302+
if self.done {
2303+
None
2304+
} else {
2305+
let nth = self.iter.nth(n);
2306+
self.done = nth.is_none();
2307+
nth
2308+
}
2309+
}
2310+
2311+
#[inline]
2312+
fn last(self) -> Option<I::Item> {
2313+
if self.done {
2314+
None
2315+
} else {
2316+
self.iter.last()
2317+
}
2318+
}
2319+
2320+
#[inline]
2321+
fn count(self) -> usize {
2322+
if self.done {
2323+
0
2324+
} else {
2325+
self.iter.count()
2326+
}
2327+
}
2328+
22022329
#[inline]
22032330
fn size_hint(&self) -> (usize, Option<usize>) {
22042331
if self.done {

0 commit comments

Comments
 (0)