Auto merge of #25035 - Stebalien:iter, r=aturon

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.

This covers:

* core::iter::Chain: count, last, nth
* core::iter::Enumerate: count, nth
* core::iter::Peekable: count, last, nth
* core::iter::Skip: count, last, next (should call nth), nth
* core::iter::Take: nth
* core::iter::Fuse: count, last, nth

of #24214.

1 files changed, 145 insertions, 18 deletions

diff --git a/src/libcore/iter.rs b/src/libcore/iter.rs
index be66623d0e1..37cabe60216 100644
--- a/src/libcore/iter.rs
+++ b/src/libcore/iter.rs
@@ -1473,6 +1473,32 @@ impl<A, B> Iterator for Chain<A, B> where
     }
 
     #[inline]
+    fn count(self) -> usize {
+        (if !self.flag { self.a.count() } else { 0 }) + self.b.count()
+    }
+
+    #[inline]
+    fn nth(&mut self, mut n: usize) -> Option<A::Item> {
+        if !self.flag {
+            for x in self.a.by_ref() {
+                if n == 0 {
+                    return Some(x)
+                }
+                n -= 1;
+            }
+            self.flag = true;
+        }
+        self.b.nth(n)
+    }
+
+    #[inline]
+    fn last(self) -> Option<A::Item> {
+        let a_last = if self.flag { None } else { self.a.last() };
+        let b_last = self.b.last();
+        b_last.or(a_last)
+    }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         let (a_lower, a_upper) = self.a.size_hint();
         let (b_lower, b_upper) = self.b.size_hint();
@@ -1777,6 +1803,20 @@ impl<I> Iterator for Enumerate<I> where I: Iterator {
     fn size_hint(&self) -> (usize, Option<usize>) {
         self.iter.size_hint()
     }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<(usize, I::Item)> {
+        self.iter.nth(n).map(|a| {
+            let i = self.count + n;
+            self.count = i + 1;
+            (i, a)
+        })
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.count()
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -1835,6 +1875,28 @@ impl<I: Iterator> Iterator for Peekable<I> {
     }
 
     #[inline]
+    fn count(self) -> usize {
+        (if self.peeked.is_some() { 1 } else { 0 }) + self.iter.count()
+    }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<I::Item> {
+        match self.peeked {
+            Some(_) if n == 0 => self.peeked.take(),
+            Some(_) => {
+                self.peeked = None;
+                self.iter.nth(n-1)
+            },
+            None => self.iter.nth(n)
+        }
+    }
+
+    #[inline]
+    fn last(self) -> Option<I::Item> {
+        self.iter.last().or(self.peeked)
+    }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         let (lo, hi) = self.iter.size_hint();
         if self.peeked.is_some() {
@@ -1960,27 +2022,49 @@ impl<I> Iterator for Skip<I> where I: Iterator {
     type Item = <I as Iterator>::Item;
 
     #[inline]
-    fn next(&mut self) -> Option<<I as Iterator>::Item> {
-        let mut next = self.iter.next();
+    fn next(&mut self) -> Option<I::Item> {
         if self.n == 0 {
-            next
+            self.iter.next()
         } else {
-            let mut n = self.n;
-            while n > 0 {
-                n -= 1;
-                match next {
-                    Some(_) => {
-                        next = self.iter.next();
-                        continue
-                    }
-                    None => {
-                        self.n = 0;
-                        return None
-                    }
-                }
-            }
+            let old_n = self.n;
             self.n = 0;
-            next
+            self.iter.nth(old_n)
+        }
+    }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<I::Item> {
+        // Can't just add n + self.n due to overflow.
+        if self.n == 0 {
+            self.iter.nth(n)
+        } else {
+            let to_skip = self.n;
+            self.n = 0;
+            // nth(n) skips n+1
+            if self.iter.nth(to_skip-1).is_none() {
+                return None;
+            }
+            self.iter.nth(n)
+        }
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        self.iter.count().saturating_sub(self.n)
+    }
+
+    #[inline]
+    fn last(mut self) -> Option<I::Item> {
+        if self.n == 0 {
+            self.iter.last()
+        } else {
+            let next = self.next();
+            if next.is_some() {
+                // recurse. n should be 0.
+                self.last().or(next)
+            } else {
+                None
+            }
         }
     }
 
@@ -2039,6 +2123,20 @@ impl<I> Iterator for Take<I> where I: Iterator{
     }
 
     #[inline]
+    fn nth(&mut self, n: usize) -> Option<I::Item> {
+        if self.n > n {
+            self.n -= n + 1;
+            self.iter.nth(n)
+        } else {
+            if self.n > 0 {
+                self.iter.nth(self.n - 1);
+                self.n = 0;
+            }
+            None
+        }
+    }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         let (lower, upper) = self.iter.size_hint();
 
@@ -2200,6 +2298,35 @@ impl<I> Iterator for Fuse<I> where I: Iterator {
     }
 
     #[inline]
+    fn nth(&mut self, n: usize) -> Option<I::Item> {
+        if self.done {
+            None
+        } else {
+            let nth = self.iter.nth(n);
+            self.done = nth.is_none();
+            nth
+        }
+    }
+
+    #[inline]
+    fn last(self) -> Option<I::Item> {
+        if self.done {
+            None
+        } else {
+            self.iter.last()
+        }
+    }
+
+    #[inline]
+    fn count(self) -> usize {
+        if self.done {
+            0
+        } else {
+            self.iter.count()
+        }
+    }
+
+    #[inline]
     fn size_hint(&self) -> (usize, Option<usize>) {
         if self.done {
             (0, Some(0))