diff --git a/src/peek_nth.rs b/src/peek_nth.rs index 1ae428ba0..5f7fb9396 100644 --- a/src/peek_nth.rs +++ b/src/peek_nth.rs @@ -115,6 +115,27 @@ where self.buf.get_mut(n) } + + /// Works exactly like the `next_if` method in `std::iter::Peekable` + pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option { + match self.next() { + Some(item) if func(&item) => Some(item), + Some(item) => { + self.buf.push_front(item); + None + } + _ => None, + } + } + + /// Works exactly like the `next_if_eq` method in `std::iter::Peekable` + pub fn next_if_eq(&mut self, expected: &T) -> Option + where + T: ?Sized, + I::Item: PartialEq, + { + self.next_if(|next| next == expected) + } } impl Iterator for PeekNth diff --git a/tests/quick.rs b/tests/quick.rs index 0d7cd9307..f0c8f2d9e 100644 --- a/tests/quick.rs +++ b/tests/quick.rs @@ -782,6 +782,30 @@ quickcheck! { assert_eq!(it.next(), None); assert_eq!(it.next(), None); } + + fn peek_nth_next_if(a: Vec) -> () { + let mut it = peek_nth(a.clone()); + for (idx, mut value) in a.iter().copied().enumerate() { + let should_be_none = it.next_if(|x| x != &value); + assert_eq!(should_be_none, None); + if value % 5 == 0 { + // Sometimes, peek up to 3 further. + let n = value as usize % 3; + let nth = it.peek_nth(n); + assert_eq!(nth, a.get(idx + n)); + } else if value % 5 == 1 { + // Sometimes, peek next element mutably. + if let Some(v) = it.peek_mut() { + *v = v.wrapping_sub(1); + let should_be_none = it.next_if_eq(&value); + assert_eq!(should_be_none, None); + value = value.wrapping_sub(1); + } + } + let eq = it.next_if_eq(&value); + assert_eq!(eq, Some(value)); + } + } } quickcheck! { diff --git a/tests/test_std.rs b/tests/test_std.rs index 3ef90f39d..732790ef4 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -738,6 +738,35 @@ fn test_peek_nth_peeking_next() { assert_eq!(iter.peek(), None); } +#[test] +fn test_peek_nth_next_if() { + let nums = vec![1u8, 2, 3, 4, 5, 6, 7]; + let mut iter = peek_nth(nums.iter().copied()); + + assert_eq!(iter.next_if(|&x| x != 0), Some(1)); + assert_eq!(iter.next(), Some(2)); + + assert_eq!(iter.peek_nth(0), Some(&3)); + assert_eq!(iter.peek_nth(1), Some(&4)); + assert_eq!(iter.next_if_eq(&3), Some(3)); + assert_eq!(iter.peek(), Some(&4)); + + assert_eq!(iter.next_if(|&x| x != 4), None); + assert_eq!(iter.next_if_eq(&4), Some(4)); + assert_eq!(iter.peek_nth(0), Some(&5)); + assert_eq!(iter.peek_nth(1), Some(&6)); + + assert_eq!(iter.next_if(|&x| x != 5), None); + assert_eq!(iter.peek(), Some(&5)); + + assert_eq!(iter.next_if(|&x| x % 2 == 1), Some(5)); + assert_eq!(iter.next_if_eq(&6), Some(6)); + assert_eq!(iter.peek_nth(0), Some(&7)); + assert_eq!(iter.peek_nth(1), None); + assert_eq!(iter.next(), Some(7)); + assert_eq!(iter.peek(), None); +} + #[test] fn pad_using() { it::assert_equal((0..0).pad_using(1, |_| 1), 1..2);