From 7b6ad606723d29b7b7f7316a019cb55c7ee48b5c Mon Sep 17 00:00:00 2001 From: Andrea Corradi Date: Fri, 3 May 2019 00:01:41 +0200 Subject: [PATCH] Add custom nth_back for Chain --- src/libcore/iter/adapters/chain.rs | 23 +++++++++++++++++++++++ src/libcore/tests/iter.rs | 16 ++++++++++++++++ 2 files changed, 39 insertions(+) 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 @@ fn next_back(&mut self) -> Option { } } + #[inline] + fn nth_back(&mut self, mut n: usize) -> Option { + 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(&mut self, init: Acc, mut f: F) -> R where Self: Sized, F: FnMut(Acc, Self::Item) -> R, R: Try { 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 @@ -103,6 +103,22 @@ fn test_iterator_chain_nth() { assert_eq!(it.next(), None); } +#[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];