about summary refs log tree commit diff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--src/libcore/iter.rs163
-rw-r--r--src/libcoretest/iter.rs215
2 files changed, 360 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))
diff --git a/src/libcoretest/iter.rs b/src/libcoretest/iter.rs
index 95a6236e9c3..0415c75aa52 100644
--- a/src/libcoretest/iter.rs
+++ b/src/libcoretest/iter.rs
@@ -101,6 +101,42 @@ fn test_iterator_chain() {
 }
 
 #[test]
+fn test_iterator_chain_nth() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let ys = [30, 40, 50, 60];
+    let zs = [];
+    let expected = [0, 1, 2, 3, 4, 5, 30, 40, 50, 60];
+    for (i, x) in expected.iter().enumerate() {
+        assert_eq!(Some(x), xs.iter().chain(ys.iter()).nth(i));
+    }
+    assert_eq!(zs.iter().chain(xs.iter()).nth(0), Some(&0));
+
+    let mut it = xs.iter().chain(zs.iter());
+    assert_eq!(it.nth(5), Some(&5));
+    assert_eq!(it.next(), None);
+}
+
+#[test]
+fn test_iterator_chain_last() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let ys = [30, 40, 50, 60];
+    let zs = [];
+    assert_eq!(xs.iter().chain(ys.iter()).last(), Some(&60));
+    assert_eq!(zs.iter().chain(ys.iter()).last(), Some(&60));
+    assert_eq!(ys.iter().chain(zs.iter()).last(), Some(&60));
+    assert_eq!(zs.iter().chain(zs.iter()).last(), None);
+}
+
+#[test]
+fn test_iterator_chain_count() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let ys = [30, 40, 50, 60];
+    let zs = [];
+    assert_eq!(xs.iter().chain(ys.iter()).count(), 10);
+    assert_eq!(zs.iter().chain(ys.iter()).count(), 4);
+}
+
+#[test]
 fn test_filter_map() {
     let it = (0..).step_by(1).take(10)
         .filter_map(|x| if x % 2 == 0 { Some(x*x) } else { None });
@@ -117,6 +153,34 @@ fn test_iterator_enumerate() {
 }
 
 #[test]
+fn test_iterator_enumerate_nth() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    for (i, &x) in xs.iter().enumerate() {
+        assert_eq!(i, x);
+    }
+
+    let mut it = xs.iter().enumerate();
+    while let Some((i, &x)) = it.nth(0) {
+        assert_eq!(i, x);
+    }
+
+    let mut it = xs.iter().enumerate();
+    while let Some((i, &x)) = it.nth(1) {
+        assert_eq!(i, x);
+    }
+
+    let (i, &x) = xs.iter().enumerate().nth(3).unwrap();
+    assert_eq!(i, x);
+    assert_eq!(i, 3);
+}
+
+#[test]
+fn test_iterator_enumerate_count() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    assert_eq!(xs.iter().count(), 6);
+}
+
+#[test]
 fn test_iterator_peekable() {
     let xs = vec![0, 1, 2, 3, 4, 5];
     let mut it = xs.iter().cloned().peekable();
@@ -149,6 +213,59 @@ fn test_iterator_peekable() {
 }
 
 #[test]
+fn test_iterator_peekable_count() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let ys = [10];
+    let zs: [i32; 0] = [];
+
+    assert_eq!(xs.iter().peekable().count(), 6);
+
+    let mut it = xs.iter().peekable();
+    assert_eq!(it.peek(), Some(&&0));
+    assert_eq!(it.count(), 6);
+
+    assert_eq!(ys.iter().peekable().count(), 1);
+
+    let mut it = ys.iter().peekable();
+    assert_eq!(it.peek(), Some(&&10));
+    assert_eq!(it.count(), 1);
+
+    assert_eq!(zs.iter().peekable().count(), 0);
+
+    let mut it = zs.iter().peekable();
+    assert_eq!(it.peek(), None);
+
+}
+
+#[test]
+fn test_iterator_peekable_nth() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let mut it = xs.iter().peekable();
+
+    assert_eq!(it.peek(), Some(&&0));
+    assert_eq!(it.nth(0), Some(&0));
+    assert_eq!(it.peek(), Some(&&1));
+    assert_eq!(it.nth(1), Some(&2));
+    assert_eq!(it.peek(), Some(&&3));
+    assert_eq!(it.nth(2), Some(&5));
+    assert_eq!(it.next(), None);
+}
+
+#[test]
+fn test_iterator_peekable_last() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let ys = [0];
+
+    let mut it = xs.iter().peekable();
+    assert_eq!(it.peek(), Some(&&0));
+    assert_eq!(it.last(), Some(&5));
+
+    let mut it = ys.iter().peekable();
+    assert_eq!(it.peek(), Some(&&0));
+    assert_eq!(it.last(), Some(&0));
+}
+
+#[test]
 fn test_iterator_take_while() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
     let ys = [0, 1, 2, 3, 5, 13];
@@ -190,6 +307,49 @@ fn test_iterator_skip() {
 }
 
 #[test]
+fn test_iterator_skip_nth() {
+    let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
+
+    let mut it = xs.iter().skip(0);
+    assert_eq!(it.nth(0), Some(&0));
+    assert_eq!(it.nth(1), Some(&2));
+
+    let mut it = xs.iter().skip(5);
+    assert_eq!(it.nth(0), Some(&13));
+    assert_eq!(it.nth(1), Some(&16));
+
+    let mut it = xs.iter().skip(12);
+    assert_eq!(it.nth(0), None);
+
+}
+
+#[test]
+fn test_iterator_skip_count() {
+    let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
+
+    assert_eq!(xs.iter().skip(0).count(), 12);
+    assert_eq!(xs.iter().skip(1).count(), 11);
+    assert_eq!(xs.iter().skip(11).count(), 1);
+    assert_eq!(xs.iter().skip(12).count(), 0);
+    assert_eq!(xs.iter().skip(13).count(), 0);
+}
+
+#[test]
+fn test_iterator_skip_last() {
+    let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19, 20, 30];
+
+    assert_eq!(xs.iter().skip(0).last(), Some(&30));
+    assert_eq!(xs.iter().skip(1).last(), Some(&30));
+    assert_eq!(xs.iter().skip(11).last(), Some(&30));
+    assert_eq!(xs.iter().skip(12).last(), None);
+    assert_eq!(xs.iter().skip(13).last(), None);
+
+    let mut it = xs.iter().skip(5);
+    assert_eq!(it.next(), Some(&13));
+    assert_eq!(it.last(), Some(&30));
+}
+
+#[test]
 fn test_iterator_take() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
     let ys = [0, 1, 2, 3, 5];
@@ -206,6 +366,30 @@ fn test_iterator_take() {
 }
 
 #[test]
+fn test_iterator_take_nth() {
+    let xs = [0, 1, 2, 4, 5];
+    let mut it = xs.iter();
+    {
+        let mut take = it.by_ref().take(3);
+        let mut i = 0;
+        while let Some(&x) = take.nth(0) {
+            assert_eq!(x, i);
+            i += 1;
+        }
+    }
+    assert_eq!(it.nth(1), Some(&5));
+    assert_eq!(it.nth(0), None);
+
+    let xs = [0, 1, 2, 3, 4];
+    let mut it = xs.iter().take(7);
+    let mut i = 1;
+    while let Some(&x) = it.nth(1) {
+        assert_eq!(x, i);
+        i += 2;
+    }
+}
+
+#[test]
 fn test_iterator_take_short() {
     let xs = [0, 1, 2, 3];
     let ys = [0, 1, 2, 3];
@@ -881,6 +1065,37 @@ fn test_fuse() {
     assert_eq!(it.len(), 0);
 }
 
+#[test]
+fn test_fuse_nth() {
+    let xs = [0, 1, 2];
+    let mut it = xs.iter();
+
+    assert_eq!(it.len(), 3);
+    assert_eq!(it.nth(2), Some(&2));
+    assert_eq!(it.len(), 0);
+    assert_eq!(it.nth(2), None);
+    assert_eq!(it.len(), 0);
+}
+
+#[test]
+fn test_fuse_last() {
+    let xs = [0, 1, 2];
+    let it = xs.iter();
+
+    assert_eq!(it.len(), 3);
+    assert_eq!(it.last(), Some(&2));
+}
+
+#[test]
+fn test_fuse_count() {
+    let xs = [0, 1, 2];
+    let it = xs.iter();
+
+    assert_eq!(it.len(), 3);
+    assert_eq!(it.count(), 3);
+    // Can't check len now because count consumes.
+}
+
 #[bench]
 fn bench_rposition(b: &mut Bencher) {
     let it: Vec<usize> = (0..300).collect();
generated by cgit-pink 1.4.1-3-g733a5 (git 2.37.1) at 2026-07-11 04:56:56 +0000