diff --git a/src/free.rs b/src/free.rs index 8d0bcf3ea..852ae0ede 100644 --- a/src/free.rs +++ b/src/free.rs @@ -20,9 +20,7 @@ pub use crate::adaptors::{interleave, put_back}; pub use crate::kmerge_impl::kmerge; pub use crate::merge_join::{merge, merge_join_by}; #[cfg(feature = "use_alloc")] -pub use crate::multipeek_impl::multipeek; -#[cfg(feature = "use_alloc")] -pub use crate::peek_nth::peek_nth; +pub use crate::multipeek_general::{multipeek, peek_nth}; #[cfg(feature = "use_alloc")] pub use crate::put_back_n_impl::put_back_n; #[cfg(feature = "use_alloc")] diff --git a/src/lib.rs b/src/lib.rs index f4de79c50..ce918d847 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -115,10 +115,8 @@ pub mod structs { pub use crate::kmerge_impl::{KMerge, KMergeBy}; pub use crate::merge_join::{Merge, MergeBy, MergeJoinBy}; #[cfg(feature = "use_alloc")] - pub use crate::multipeek_impl::MultiPeek; + pub use crate::multipeek_general::{MultiPeek, PeekNth}; pub use crate::pad_tail::PadUsing; - #[cfg(feature = "use_alloc")] - pub use crate::peek_nth::PeekNth; pub use crate::peeking_take_while::PeekingTakeWhile; #[cfg(feature = "use_alloc")] pub use crate::permutations::Permutations; @@ -205,10 +203,8 @@ mod lazy_buffer; mod merge_join; mod minmax; #[cfg(feature = "use_alloc")] -mod multipeek_impl; +mod multipeek_general; mod pad_tail; -#[cfg(feature = "use_alloc")] -mod peek_nth; mod peeking_take_while; #[cfg(feature = "use_alloc")] mod permutations; @@ -4115,7 +4111,7 @@ pub trait Itertools: Iterator { where Self: Sized, { - multipeek_impl::multipeek(self) + multipeek_general::multipeek(self) } /// Collect the items in this iterator and return a `HashMap` which diff --git a/src/peek_nth.rs b/src/multipeek_general.rs similarity index 55% rename from src/peek_nth.rs rename to src/multipeek_general.rs index b03a3ef5f..6759b735f 100644 --- a/src/peek_nth.rs +++ b/src/multipeek_general.rs @@ -1,17 +1,34 @@ -use crate::size_hint; -use crate::PeekingNext; +use crate::{size_hint, PeekingNext}; use alloc::collections::VecDeque; use std::iter::Fuse; -/// See [`peek_nth()`] for more information. -#[derive(Clone, Debug)] #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] -pub struct PeekNth -where - I: Iterator, -{ +#[derive(Debug, Clone)] +pub struct MultiPeekGeneral { iter: Fuse, buf: VecDeque, + index: Idx, +} + +/// See [`multipeek()`] for more information. +pub type MultiPeek = MultiPeekGeneral; + +/// See [`peek_nth()`] for more information. +pub type PeekNth = MultiPeekGeneral; + +/// An iterator adaptor that allows the user to peek at multiple `.next()` +/// values without advancing the base iterator. +/// +/// [`IntoIterator`] enabled version of [`Itertools::multipeek`](crate::Itertools::multipeek). +pub fn multipeek(iterable: I) -> MultiPeek +where + I: IntoIterator, +{ + MultiPeek { + iter: iterable.into_iter().fuse(), + buf: VecDeque::new(), + index: 0, + } } /// A drop-in replacement for [`std::iter::Peekable`] which adds a `peek_nth` @@ -28,16 +45,60 @@ where PeekNth { iter: iterable.into_iter().fuse(), buf: VecDeque::new(), + index: (), } } -impl PeekNth -where - I: Iterator, -{ - /// Works exactly like the `peek` method in [`std::iter::Peekable`]. +pub trait PeekIndex { + fn reset_index(&mut self); + fn add(&mut self, n: usize); + fn index(&self) -> usize; +} + +impl PeekIndex for () { + fn reset_index(&mut self) {} + fn add(&mut self, _n: usize) {} + fn index(&self) -> usize { + 0 + } +} + +impl PeekIndex for usize { + fn reset_index(&mut self) { + *self = 0; + } + fn add(&mut self, n: usize) { + *self += n + } + fn index(&self) -> usize { + *self + } +} + +impl MultiPeekGeneral { + /// Works similarly to `.next()`, but does not advance the iterator itself. + /// + /// - For `MultiPeekGeneral`, calling `.peek()` will increment the internal index, + /// so the next call to `.peek()` will advance to the next element in the buffer. + /// The actual underlying iterator is not consumed, allowing multiple peeks + /// without advancing the iterator itself. + /// - For `peek_nth`, since there is no internal index used, calling `.peek()` + /// multiple times will not advance the internal state or the iterator, providing + /// a consistent view of the same element. pub fn peek(&mut self) -> Option<&I::Item> { - self.peek_nth(0) + let ret = if self.index.index() < self.buf.len() { + Some(&self.buf[self.index.index()]) + } else { + match self.iter.next() { + Some(x) => { + self.buf.push_back(x); + Some(&self.buf[self.index.index()]) + } + None => return None, + } + }; + self.index.add(1); + ret } /// Works exactly like the `peek_mut` method in [`std::iter::Peekable`]. @@ -73,7 +134,13 @@ where 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. @@ -110,6 +177,8 @@ where /// 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)); @@ -137,15 +206,36 @@ where { self.next_if(|next| next == expected) } + + /// Works exactly like `next_if`, but for the `nth` value without advancing the iterator. + pub fn nth_if(&mut self, n: usize, func: impl FnOnce(&I::Item) -> bool) -> Option<&I::Item> { + match self.peek_nth(n) { + Some(item) if func(item) => Some(item), + _ => None, + } + } + + /// Works exactly like `next_if_eq`, but for the `nth` value without advancing the iterator. + pub fn nth_if_eq(&mut self, n: usize, expected: &T) -> Option<&I::Item> + where + T: ?Sized, + I::Item: PartialEq, + { + self.nth_if(n, |item| item == expected) + } +} +impl MultiPeek { + /// Reset the peeking “cursor” + pub fn reset_peek(&mut self) { + self.index = 0 + } } -impl Iterator for PeekNth -where - I: Iterator, -{ +impl Iterator for MultiPeekGeneral { type Item = I::Item; fn next(&mut self) -> Option { + self.index.reset_index(); self.buf.pop_front().or_else(|| self.iter.next()) } @@ -162,17 +252,28 @@ where } } -impl ExactSizeIterator for PeekNth where I: ExactSizeIterator {} +impl ExactSizeIterator for MultiPeekGeneral {} -impl PeekingNext for PeekNth +impl PeekingNext for MultiPeekGeneral where I: Iterator, + Idx: PeekIndex, { fn peeking_next(&mut self, accept: F) -> Option where F: FnOnce(&Self::Item) -> bool, { - self.peek().filter(|item| accept(item))?; + if self.buf.is_empty() { + if let Some(r) = self.peek() { + if !accept(r) { + return None; + } + } + } else if let Some(r) = self.buf.front() { + if !accept(r) { + return None; + } + } self.next() } } diff --git a/src/multipeek_impl.rs b/src/multipeek_impl.rs deleted file mode 100644 index 6f800b6fb..000000000 --- a/src/multipeek_impl.rs +++ /dev/null @@ -1,116 +0,0 @@ -use crate::size_hint; -#[cfg(doc)] -use crate::Itertools; -use crate::PeekingNext; -use alloc::collections::VecDeque; -use std::iter::Fuse; - -/// See [`multipeek()`] for more information. -#[derive(Clone, Debug)] -#[must_use = "iterator adaptors are lazy and do nothing unless consumed"] -pub struct MultiPeek -where - I: Iterator, -{ - iter: Fuse, - buf: VecDeque, - index: usize, -} - -/// An iterator adaptor that allows the user to peek at multiple `.next()` -/// values without advancing the base iterator. -/// -/// [`IntoIterator`] enabled version of [`Itertools::multipeek`]. -pub fn multipeek(iterable: I) -> MultiPeek -where - I: IntoIterator, -{ - MultiPeek { - iter: iterable.into_iter().fuse(), - buf: VecDeque::new(), - index: 0, - } -} - -impl MultiPeek -where - I: Iterator, -{ - /// Reset the peeking “cursor” - pub fn reset_peek(&mut self) { - self.index = 0; - } -} - -impl MultiPeek { - /// Works exactly like `.next()` with the only difference that it doesn't - /// advance itself. `.peek()` can be called multiple times, to peek - /// further ahead. - /// When `.next()` is called, reset the peeking “cursor”. - pub fn peek(&mut self) -> Option<&I::Item> { - let ret = if self.index < self.buf.len() { - Some(&self.buf[self.index]) - } else { - match self.iter.next() { - Some(x) => { - self.buf.push_back(x); - Some(&self.buf[self.index]) - } - None => return None, - } - }; - - self.index += 1; - ret - } -} - -impl PeekingNext for MultiPeek -where - I: Iterator, -{ - fn peeking_next(&mut self, accept: F) -> Option - where - F: FnOnce(&Self::Item) -> bool, - { - if self.buf.is_empty() { - if let Some(r) = self.peek() { - if !accept(r) { - return None; - } - } - } else if let Some(r) = self.buf.front() { - if !accept(r) { - return None; - } - } - self.next() - } -} - -impl Iterator for MultiPeek -where - I: Iterator, -{ - type Item = I::Item; - - fn next(&mut self) -> Option { - self.index = 0; - self.buf.pop_front().or_else(|| self.iter.next()) - } - - fn size_hint(&self) -> (usize, Option) { - size_hint::add_scalar(self.iter.size_hint(), self.buf.len()) - } - - fn fold(self, mut init: B, mut f: F) -> B - where - F: FnMut(B, Self::Item) -> B, - { - init = self.buf.into_iter().fold(init, &mut f); - self.iter.fold(init, f) - } -} - -// Same size -impl ExactSizeIterator for MultiPeek where I: ExactSizeIterator {} 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];