Add custom nth_back for Chain

2 files changed, 39 insertions, 0 deletions

diff --git a/src/libcore/iter/adapters/chain.rs b/src/libcore/iter/adapters/chain.rs
index 76239ebc0ab..0b9f7f6b609 100644
--- a/src/libcore/iter/adapters/chain.rs
+++ b/src/libcore/iter/adapters/chain.rs
@@ -207,6 +207,29 @@ impl<A, B> DoubleEndedIterator for Chain<A, B> where
         }
     }
 
+    #[inline]
+    fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
+        match self.state {
+            ChainState::Both | ChainState::Back => {
+                for x in self.b.by_ref().rev() {
+                    if n == 0 {
+                        return Some(x)
+                    }
+                    n -= 1;
+                }
+                if let ChainState::Both = self.state {
+                    self.state = ChainState::Front;
+                }
+            }
+            ChainState::Front => {}
+        }
+        if let ChainState::Front = self.state {
+            self.a.nth_back(n)
+        } else {
+            None
+        }
+    }
+
     fn try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R where
         Self: Sized, F: FnMut(Acc, Self::Item) -> R, R: Try<Ok=Acc>
     {
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 7dfb1adad9e..449e339f89c 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -104,6 +104,22 @@ fn test_iterator_chain_nth() {
 }
 
 #[test]
+fn test_iterator_chain_nth_back() {
+    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().rev().enumerate() {
+        assert_eq!(Some(x), xs.iter().chain(&ys).nth_back(i));
+    }
+    assert_eq!(zs.iter().chain(&xs).nth_back(0), Some(&5));
+
+    let mut it = xs.iter().chain(&zs);
+    assert_eq!(it.nth_back(5), Some(&0));
+    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];