diff --git a/src/multipeek_general.rs b/src/multipeek_general.rs index 64675d2d6..6759b735f 100644 --- a/src/multipeek_general.rs +++ b/src/multipeek_general.rs @@ -51,13 +51,13 @@ where pub trait PeekIndex { fn reset_index(&mut self); - fn increment_index(&mut self); + fn add(&mut self, n: usize); fn index(&self) -> usize; } impl PeekIndex for () { fn reset_index(&mut self) {} - fn increment_index(&mut self) {} + fn add(&mut self, _n: usize) {} fn index(&self) -> usize { 0 } @@ -67,8 +67,8 @@ impl PeekIndex for usize { fn reset_index(&mut self) { *self = 0; } - fn increment_index(&mut self) { - *self += 1 + fn add(&mut self, n: usize) { + *self += n } fn index(&self) -> usize { *self @@ -97,8 +97,7 @@ impl MultiPeekGeneral { None => return None, } }; - - self.index.increment_index(); + self.index.add(1); ret } @@ -135,7 +134,13 @@ impl MultiPeekGeneral { self.buf.extend(self.iter.by_ref().take(unbuffered_items)); - self.buf.get(n) + let ret = self.buf.get(n); + + if ret.is_some() { + self.index.add(n + 1); + } + + ret } /// Returns a mutable reference to the `nth` value without advancing the iterator. @@ -172,6 +177,8 @@ impl MultiPeekGeneral { /// assert_eq!(iter.peek_nth_mut(1), None); /// ``` pub fn peek_nth_mut(&mut self, n: usize) -> Option<&mut I::Item> { + self.index.add(n); + let unbuffered_items = (n + 1).saturating_sub(self.buf.len()); self.buf.extend(self.iter.by_ref().take(unbuffered_items)); diff --git a/tests/test_std.rs b/tests/test_std.rs index 00246d506..ef55f190d 100644 --- a/tests/test_std.rs +++ b/tests/test_std.rs @@ -660,6 +660,20 @@ fn test_multipeek() { assert_eq!(mp.peek(), None); } +#[test] +fn test_multipeek_peeknth() { + let nums = vec![6, 7, 8, 9, 10]; + + let mut mp = multipeek(nums); + assert_eq!(mp.peek_nth(2), Some(&8)); + assert_eq!(mp.peek(), Some(&9)); + assert_eq!(mp.peek(), Some(&10)); + mp.reset_peek(); + assert_eq!(mp.peek_nth(1), Some(&7)); + assert_eq!(mp.peek_nth(3), Some(&9)); + assert_eq!(mp.peek_nth(10), None); +} + #[test] fn test_multipeek_reset() { let data = [1, 2, 3, 4];