Implement DoubleEndedIterator for iter::{StepBy, Peekable, Take}

2 files changed, 296 insertions, 13 deletions

diff --git a/src/libcore/iter/adapters/mod.rs b/src/libcore/iter/adapters/mod.rs
index 518442efe74..00053d9da84 100644
--- a/src/libcore/iter/adapters/mod.rs
+++ b/src/libcore/iter/adapters/mod.rs
@@ -485,6 +485,39 @@ impl<I> Iterator for StepBy<I> where I: Iterator {
     }
 }
 
+impl<I> StepBy<I> where I: ExactSizeIterator {
+    // The zero-based index starting from the end of the iterator of the
+    // last element. Used in the `DoubleEndedIterator` implementation.
+    fn next_back_index(&self) -> usize {
+        let rem = self.iter.len() % (self.step + 1);
+        if self.first_take {
+            if rem == 0 { self.step } else { rem - 1 }
+        } else {
+            rem
+        }
+    }
+}
+
+#[stable(feature = "double_ended_step_by_iterator", since = "1.38.0")]
+impl<I> DoubleEndedIterator for StepBy<I> where I: DoubleEndedIterator + ExactSizeIterator {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.nth_back(self.next_back_index())
+    }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        // `self.iter.nth_back(usize::MAX)` does the right thing here when `n`
+        // is out of bounds because the length of `self.iter` does not exceed
+        // `usize::MAX` (because `I: ExactSizeIterator`) and `nth_back` is
+        // zero-indexed
+        let n = n
+            .saturating_mul(self.step + 1)
+            .saturating_add(self.next_back_index());
+        self.iter.nth_back(n)
+    }
+}
+
 // StepBy can only make the iterator shorter, so the len will still fit.
 #[stable(feature = "iterator_step_by", since = "1.28.0")]
 impl<I> ExactSizeIterator for StepBy<I> where I: ExactSizeIterator {}
@@ -1158,6 +1191,45 @@ impl<I: Iterator> Iterator for Peekable<I> {
     }
 }
 
+#[stable(feature = "double_ended_peek_iterator", since = "1.38.0")]
+impl<I> DoubleEndedIterator for Peekable<I> where I: DoubleEndedIterator {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        self.iter.next_back().or_else(|| self.peeked.take().and_then(|x| x))
+    }
+
+    #[inline]
+    fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R where
+        Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
+    {
+        match self.peeked.take() {
+            Some(None) => return Try::from_ok(init),
+            Some(Some(v)) => match self.iter.try_rfold(init, &mut f).into_result() {
+                Ok(acc) => f(acc, v),
+                Err(e) => {
+                    self.peeked = Some(Some(v));
+                    Try::from_error(e)
+                }
+            },
+            None => self.iter.try_rfold(init, f),
+        }
+    }
+
+    #[inline]
+    fn rfold<Acc, Fold>(self, init: Acc, mut fold: Fold) -> Acc
+        where Fold: FnMut(Acc, Self::Item) -> Acc,
+    {
+        match self.peeked {
+            Some(None) => return init,
+            Some(Some(v)) => {
+                let acc = self.iter.rfold(init, &mut fold);
+                fold(acc, v)
+            }
+            None => self.iter.rfold(init, fold),
+        }
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<I: ExactSizeIterator> ExactSizeIterator for Peekable<I> {}
 
@@ -1613,6 +1685,51 @@ impl<I> Iterator for Take<I> where I: Iterator{
     }
 }
 
+#[stable(feature = "double_ended_take_iterator", since = "1.38.0")]
+impl<I> DoubleEndedIterator for Take<I> where I: DoubleEndedIterator + ExactSizeIterator {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.n == 0 {
+            None
+        } else {
+            let n = self.n;
+            self.n -= 1;
+            self.iter.nth_back(self.iter.len().saturating_sub(n))
+        }
+    }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        let len = self.iter.len();
+        if self.n > n {
+            let m = len.saturating_sub(self.n) + n;
+            self.n -= n + 1;
+            self.iter.nth_back(m)
+        } else {
+            if len > 0 {
+                self.iter.nth_back(len - 1);
+            }
+            None
+        }
+    }
+
+    #[inline]
+    fn try_rfold<Acc, Fold, R>(&mut self, init: Acc, fold: Fold) -> R where
+        Self: Sized, Fold: FnMut(Acc, Self::Item) -> R, R: Try<Ok = Acc>
+    {
+        if self.n == 0 {
+            Try::from_ok(init)
+        } else {
+            let len = self.iter.len();
+            if len > self.n && self.iter.nth_back(len - self.n - 1).is_none() {
+                Try::from_ok(init)
+            } else {
+                self.iter.try_rfold(init, fold)
+            }
+        }
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<I> ExactSizeIterator for Take<I> where I: ExactSizeIterator {}
 
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index bedb9e75612..60efeda3a40 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -188,6 +188,19 @@ fn test_iterator_step_by() {
     assert_eq!(it.next(), Some(6));
     assert_eq!(it.next(), Some(9));
     assert_eq!(it.next(), None);
+
+    let mut it = (0..3).step_by(1);
+    assert_eq!(it.next_back(), Some(2));
+    assert_eq!(it.next_back(), Some(1));
+    assert_eq!(it.next_back(), Some(0));
+    assert_eq!(it.next_back(), None);
+
+    let mut it = (0..11).step_by(3);
+    assert_eq!(it.next_back(), Some(9));
+    assert_eq!(it.next_back(), Some(6));
+    assert_eq!(it.next_back(), Some(3));
+    assert_eq!(it.next_back(), Some(0));
+    assert_eq!(it.next_back(), None);
 }
 
 #[test]
@@ -253,6 +266,31 @@ fn test_iterator_step_by_nth_overflow() {
 }
 
 #[test]
+fn test_iterator_step_by_nth_back() {
+    let mut it = (0..16).step_by(5);
+    assert_eq!(it.nth_back(0), Some(15));
+    assert_eq!(it.nth_back(0), Some(10));
+    assert_eq!(it.nth_back(0), Some(5));
+    assert_eq!(it.nth_back(0), Some(0));
+    assert_eq!(it.nth_back(0), None);
+
+    let mut it = (0..16).step_by(5);
+    assert_eq!(it.next(), Some(0)); // to set `first_take` to `false`
+    assert_eq!(it.nth_back(0), Some(15));
+    assert_eq!(it.nth_back(0), Some(10));
+    assert_eq!(it.nth_back(0), Some(5));
+    assert_eq!(it.nth_back(0), None);
+
+    let it = || (0..18).step_by(5);
+    assert_eq!(it().nth_back(0), Some(15));
+    assert_eq!(it().nth_back(1), Some(10));
+    assert_eq!(it().nth_back(2), Some(5));
+    assert_eq!(it().nth_back(3), Some(0));
+    assert_eq!(it().nth_back(4), None);
+    assert_eq!(it().nth_back(42), None);
+}
+
+#[test]
 #[should_panic]
 fn test_iterator_step_by_zero() {
     let mut it = (0..).step_by(0);
@@ -465,8 +503,8 @@ fn test_iterator_filter_fold() {
 #[test]
 fn test_iterator_peekable() {
     let xs = vec![0, 1, 2, 3, 4, 5];
-    let mut it = xs.iter().cloned().peekable();
 
+    let mut it = xs.iter().cloned().peekable();
     assert_eq!(it.len(), 6);
     assert_eq!(it.peek().unwrap(), &0);
     assert_eq!(it.len(), 6);
@@ -492,6 +530,33 @@ fn test_iterator_peekable() {
     assert_eq!(it.len(), 0);
     assert!(it.next().is_none());
     assert_eq!(it.len(), 0);
+
+    let mut it = xs.iter().cloned().peekable();
+    assert_eq!(it.len(), 6);
+    assert_eq!(it.peek().unwrap(), &0);
+    assert_eq!(it.len(), 6);
+    assert_eq!(it.next_back().unwrap(), 5);
+    assert_eq!(it.len(), 5);
+    assert_eq!(it.next_back().unwrap(), 4);
+    assert_eq!(it.len(), 4);
+    assert_eq!(it.next_back().unwrap(), 3);
+    assert_eq!(it.len(), 3);
+    assert_eq!(it.peek().unwrap(), &0);
+    assert_eq!(it.len(), 3);
+    assert_eq!(it.peek().unwrap(), &0);
+    assert_eq!(it.len(), 3);
+    assert_eq!(it.next_back().unwrap(), 2);
+    assert_eq!(it.len(), 2);
+    assert_eq!(it.next_back().unwrap(), 1);
+    assert_eq!(it.len(), 1);
+    assert_eq!(it.peek().unwrap(), &0);
+    assert_eq!(it.len(), 1);
+    assert_eq!(it.next_back().unwrap(), 0);
+    assert_eq!(it.len(), 0);
+    assert!(it.peek().is_none());
+    assert_eq!(it.len(), 0);
+    assert!(it.next_back().is_none());
+    assert_eq!(it.len(), 0);
 }
 
 #[test]
@@ -564,6 +629,18 @@ fn test_iterator_peekable_fold() {
     assert_eq!(i, xs.len());
 }
 
+#[test]
+fn test_iterator_peekable_rfold() {
+    let xs = [0, 1, 2, 3, 4, 5];
+    let mut it = xs.iter().peekable();
+    assert_eq!(it.peek(), Some(&&0));
+    let i = it.rfold(0, |i, &x| {
+        assert_eq!(x, xs[xs.len() - 1 - i]);
+        i + 1
+    });
+    assert_eq!(i, xs.len());
+}
+
 /// This is an iterator that follows the Iterator contract,
 /// but it is not fused. After having returned None once, it will start
 /// producing elements if .next() is called again.
@@ -812,13 +889,25 @@ fn test_iterator_skip_fold() {
 fn test_iterator_take() {
     let xs = [0, 1, 2, 3, 5, 13, 15, 16, 17, 19];
     let ys = [0, 1, 2, 3, 5];
-    let mut it = xs.iter().take(5);
+
+    let mut it = xs.iter().take(ys.len());
     let mut i = 0;
-    assert_eq!(it.len(), 5);
+    assert_eq!(it.len(), ys.len());
     while let Some(&x) = it.next() {
         assert_eq!(x, ys[i]);
         i += 1;
-        assert_eq!(it.len(), 5-i);
+        assert_eq!(it.len(), ys.len() - i);
+    }
+    assert_eq!(i, ys.len());
+    assert_eq!(it.len(), 0);
+
+    let mut it = xs.iter().take(ys.len());
+    let mut i = 0;
+    assert_eq!(it.len(), ys.len());
+    while let Some(&x) = it.next_back() {
+        i += 1;
+        assert_eq!(x, ys[ys.len() - i]);
+        assert_eq!(it.len(), ys.len() - i);
     }
     assert_eq!(i, ys.len());
     assert_eq!(it.len(), 0);
@@ -849,18 +938,50 @@ fn test_iterator_take_nth() {
 }
 
 #[test]
+fn test_iterator_take_nth_back() {
+    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_back(0) {
+            i += 1;
+            assert_eq!(x, 3 - i);
+        }
+    }
+    assert_eq!(it.nth_back(0), None);
+
+    let xs = [0, 1, 2, 3, 4];
+    let mut it = xs.iter().take(7);
+    assert_eq!(it.nth_back(1), Some(&3));
+    assert_eq!(it.nth_back(1), Some(&1));
+    assert_eq!(it.nth_back(1), None);
+}
+
+#[test]
 fn test_iterator_take_short() {
     let xs = [0, 1, 2, 3];
-    let ys = [0, 1, 2, 3];
+
     let mut it = xs.iter().take(5);
     let mut i = 0;
-    assert_eq!(it.len(), 4);
+    assert_eq!(it.len(), xs.len());
     while let Some(&x) = it.next() {
-        assert_eq!(x, ys[i]);
+        assert_eq!(x, xs[i]);
         i += 1;
-        assert_eq!(it.len(), 4-i);
+        assert_eq!(it.len(), xs.len() - i);
     }
-    assert_eq!(i, ys.len());
+    assert_eq!(i, xs.len());
+    assert_eq!(it.len(), 0);
+
+    let mut it = xs.iter().take(5);
+    let mut i = 0;
+    assert_eq!(it.len(), xs.len());
+    while let Some(&x) = it.next_back() {
+        i += 1;
+        assert_eq!(x, xs[xs.len() - i]);
+        assert_eq!(it.len(), xs.len() - i);
+    }
+    assert_eq!(i, xs.len());
     assert_eq!(it.len(), 0);
 }
 
@@ -2241,17 +2362,50 @@ fn test_enumerate_try_folds() {
 }
 
 #[test]
-fn test_peek_try_fold() {
+fn test_peek_try_folds() {
     let f = &|acc, x| i32::checked_add(2*acc, x);
+
     assert_eq!((1..20).peekable().try_fold(7, f), (1..20).try_fold(7, f));
+    assert_eq!((1..20).peekable().try_rfold(7, f), (1..20).try_rfold(7, f));
+
     let mut iter = (1..20).peekable();
     assert_eq!(iter.peek(), Some(&1));
     assert_eq!(iter.try_fold(7, f), (1..20).try_fold(7, f));
 
+    let mut iter = (1..20).peekable();
+    assert_eq!(iter.peek(), Some(&1));
+    assert_eq!(iter.try_rfold(7, f), (1..20).try_rfold(7, f));
+
     let mut iter = [100, 20, 30, 40, 50, 60, 70].iter().cloned().peekable();
     assert_eq!(iter.peek(), Some(&100));
     assert_eq!(iter.try_fold(0, i8::checked_add), None);
     assert_eq!(iter.peek(), Some(&40));
+
+    let mut iter = [100, 20, 30, 40, 50, 60, 70].iter().cloned().peekable();
+    assert_eq!(iter.peek(), Some(&100));
+    assert_eq!(iter.try_rfold(0, i8::checked_add), None);
+    assert_eq!(iter.peek(), Some(&100));
+    assert_eq!(iter.next_back(), Some(50));
+
+    let mut iter = (2..5).peekable();
+    assert_eq!(iter.peek(), Some(&2));
+    assert_eq!(iter.try_for_each(Err), Err(2));
+    assert_eq!(iter.peek(), Some(&3));
+    assert_eq!(iter.try_for_each(Err), Err(3));
+    assert_eq!(iter.peek(), Some(&4));
+    assert_eq!(iter.try_for_each(Err), Err(4));
+    assert_eq!(iter.peek(), None);
+    assert_eq!(iter.try_for_each(Err), Ok(()));
+
+    let mut iter = (2..5).peekable();
+    assert_eq!(iter.peek(), Some(&2));
+    assert_eq!(iter.try_rfold((), |(), x| Err(x)), Err(4));
+    assert_eq!(iter.peek(), Some(&2));
+    assert_eq!(iter.try_rfold((), |(), x| Err(x)), Err(3));
+    assert_eq!(iter.peek(), Some(&2));
+    assert_eq!(iter.try_rfold((), |(), x| Err(x)), Err(2));
+    assert_eq!(iter.peek(), None);
+    assert_eq!(iter.try_rfold((), |(), x| Err(x)), Ok(()));
 }
 
 #[test]
@@ -2300,13 +2454,25 @@ fn test_skip_try_folds() {
 fn test_take_try_folds() {
     let f = &|acc, x| i32::checked_add(2*acc, x);
     assert_eq!((10..30).take(10).try_fold(7, f), (10..20).try_fold(7, f));
-    //assert_eq!((10..30).take(10).try_rfold(7, f), (10..20).try_rfold(7, f));
+    assert_eq!((10..30).take(10).try_rfold(7, f), (10..20).try_rfold(7, f));
 
     let mut iter = (10..30).take(20);
     assert_eq!(iter.try_fold(0, i8::checked_add), None);
     assert_eq!(iter.next(), Some(20));
-    //assert_eq!(iter.try_rfold(0, i8::checked_add), None);
-    //assert_eq!(iter.next_back(), Some(24));
+    assert_eq!(iter.try_rfold(0, i8::checked_add), None);
+    assert_eq!(iter.next_back(), Some(24));
+
+    let mut iter = (2..20).take(3);
+    assert_eq!(iter.try_for_each(Err), Err(2));
+    assert_eq!(iter.try_for_each(Err), Err(3));
+    assert_eq!(iter.try_for_each(Err), Err(4));
+    assert_eq!(iter.try_for_each(Err), Ok(()));
+
+    let mut iter = (2..20).take(3).rev();
+    assert_eq!(iter.try_for_each(Err), Err(4));
+    assert_eq!(iter.try_for_each(Err), Err(3));
+    assert_eq!(iter.try_for_each(Err), Err(2));
+    assert_eq!(iter.try_for_each(Err), Ok(()));
 }
 
 #[test]