Merge multipeek peeknth

src/free.rs

@@ -20,9 +20,9 @@ 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;
+pub use crate::multipeek_general::multipeek;
 #[cfg(feature = "use_alloc")]
-pub use crate::peek_nth::peek_nth;
+pub use crate::multipeek_general::peek_nth;
 #[cfg(feature = "use_alloc")]
 pub use crate::put_back_n_impl::put_back_n;
 #[cfg(feature = "use_alloc")]

src/lib.rs

@@ -115,10 +115,10 @@ 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::pad_tail::PadUsing;
+    pub use crate::multipeek_general::MultiPeek;
     #[cfg(feature = "use_alloc")]
-    pub use crate::peek_nth::PeekNth;
+    pub use crate::multipeek_general::PeekNth;
+    pub use crate::pad_tail::PadUsing;
     pub use crate::peeking_take_while::PeekingTakeWhile;
     #[cfg(feature = "use_alloc")]
     pub use crate::permutations::Permutations;
@@ -205,10 +205,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 +4113,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

src/peek_nth.rs → src/multipeek_general.rs

@@ -1,17 +1,38 @@
-use crate::size_hint;
-use crate::PeekingNext;
+#![allow(private_interfaces)]
+#![allow(private_bounds)]
+
+use crate::{size_hint, PeekingNext};
 use alloc::collections::VecDeque;
 use std::iter::Fuse;
 
-/// See [`peek_nth()`] for more information.
-#[derive(Clone, Debug)]
+/// See [`multipeek()`] for more information.
 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
-pub struct PeekNth<I>
+#[derive(Debug, Clone)]
+pub struct MultiPeekGeneral<I: Iterator, Idx> {
+    pub iter: Fuse<I>,
+    pub buf: VecDeque<I::Item>,
+    pub index: Idx,
+}
+
+/// See [`multipeek()`] for more information.
+pub type MultiPeek<I> = MultiPeekGeneral<I, usize>;
+
+/// See [`peek_nth()`] for more information.
+pub type PeekNth<I> = MultiPeekGeneral<I, ()>;
+
+/// An iterator adaptor that allows the user to peek at multiple `.next()`
+/// values without advancing the base iterator.
+///
+/// [`IntoIterator`] enabled version of [`crate::Itertools::multipeek`].
+pub fn multipeek<I>(iterable: I) -> MultiPeek<I::IntoIter>
 where
-    I: Iterator,
+    I: IntoIterator,
 {
-    iter: Fuse<I>,
-    buf: VecDeque<I::Item>,
+    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,18 +49,25 @@ where
     PeekNth {
         iter: iterable.into_iter().fuse(),
         buf: VecDeque::new(),
+        index: (),
     }
 }
 
-impl<I> PeekNth<I>
-where
-    I: Iterator,
-{
-    /// Works exactly like the `peek` method in [`std::iter::Peekable`].
-    pub fn peek(&mut self) -> Option<&I::Item> {
-        self.peek_nth(0)
+pub trait PeekIndex {
+    fn reset_index(&mut self);
+}
+
+impl PeekIndex for () {
+    fn reset_index(&mut self) {}
+}
+
+impl PeekIndex for usize {
+    fn reset_index(&mut self) {
+        *self = 0;
     }
+}
 
+impl<I: Iterator, Idx: PeekIndex> MultiPeekGeneral<I, Idx> {
     /// Works exactly like the `peek_mut` method in [`std::iter::Peekable`].
     pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
         self.peek_nth_mut(0)
@@ -137,15 +165,64 @@ 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<T>(&mut self, n: usize, expected: &T) -> Option<&I::Item>
+    where
+        T: ?Sized,
+        I::Item: PartialEq<T>,
+    {
+        self.nth_if(n, |item| item == expected)
+    }
 }
+impl<I: Iterator> MultiPeek<I> {
+    /// Reset the peeking “cursor”
+    pub fn reset_peek(&mut self) {
+        self.index = 0
+    }
 
-impl<I> Iterator for PeekNth<I>
-where
-    I: Iterator,
-{
+    /// 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<I: Iterator> PeekNth<I> {
+    /// Works exactly like the `peek` method in [`std::iter::Peekable`].
+    pub fn peek(&mut self) -> Option<&I::Item> {
+        self.peek_nth(0)
+    }
+}
+
+impl<I: Iterator, Idx: PeekIndex> Iterator for MultiPeekGeneral<I, Idx> {
     type Item = I::Item;
 
     fn next(&mut self) -> Option<Self::Item> {
+        self.index.reset_index();
         self.buf.pop_front().or_else(|| self.iter.next())
     }
 
@@ -162,17 +239,17 @@ where
     }
 }
 
-impl<I> ExactSizeIterator for PeekNth<I> where I: ExactSizeIterator {}
+impl<I: ExactSizeIterator, Idx: PeekIndex> ExactSizeIterator for MultiPeekGeneral<I, Idx> {}
 
-impl<I> PeekingNext for PeekNth<I>
+impl<I: Iterator, Idx: PeekIndex> PeekingNext for MultiPeekGeneral<I, Idx>
 where
     I: Iterator,
 {
     fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
     where
         F: FnOnce(&Self::Item) -> bool,
     {
-        self.peek().filter(|item| accept(item))?;
+        self.peek_mut().filter(|item| accept(item))?;
         self.next()
     }
 }