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

Author: timvermeulen

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 {}