itertools/
peek_nth.rs

1use crate::size_hint;
2use crate::PeekingNext;
3use alloc::collections::VecDeque;
4use std::iter::Fuse;
5
6/// See [`peek_nth()`] for more information.
7#[derive(Clone, Debug)]
8#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
9pub struct PeekNth<I>
10where
11    I: Iterator,
12{
13    iter: Fuse<I>,
14    buf: VecDeque<I::Item>,
15}
16
17/// A drop-in replacement for [`std::iter::Peekable`] which adds a `peek_nth`
18/// method allowing the user to `peek` at a value several iterations forward
19/// without advancing the base iterator.
20///
21/// This differs from `multipeek` in that subsequent calls to `peek` or
22/// `peek_nth` will always return the same value until `next` is called
23/// (making `reset_peek` unnecessary).
24pub fn peek_nth<I>(iterable: I) -> PeekNth<I::IntoIter>
25where
26    I: IntoIterator,
27{
28    PeekNth {
29        iter: iterable.into_iter().fuse(),
30        buf: VecDeque::new(),
31    }
32}
33
34impl<I> PeekNth<I>
35where
36    I: Iterator,
37{
38    /// Works exactly like the `peek` method in [`std::iter::Peekable`].
39    pub fn peek(&mut self) -> Option<&I::Item> {
40        self.peek_nth(0)
41    }
42
43    /// Works exactly like the `peek_mut` method in [`std::iter::Peekable`].
44    pub fn peek_mut(&mut self) -> Option<&mut I::Item> {
45        self.peek_nth_mut(0)
46    }
47
48    /// Returns a reference to the `nth` value without advancing the iterator.
49    ///
50    /// # Examples
51    ///
52    /// Basic usage:
53    ///
54    /// ```
55    /// use itertools::peek_nth;
56    ///
57    /// let xs = vec![1, 2, 3];
58    /// let mut iter = peek_nth(xs.into_iter());
59    ///
60    /// assert_eq!(iter.peek_nth(0), Some(&1));
61    /// assert_eq!(iter.next(), Some(1));
62    ///
63    /// // The iterator does not advance even if we call `peek_nth` multiple times
64    /// assert_eq!(iter.peek_nth(0), Some(&2));
65    /// assert_eq!(iter.peek_nth(1), Some(&3));
66    /// assert_eq!(iter.next(), Some(2));
67    ///
68    /// // Calling `peek_nth` past the end of the iterator will return `None`
69    /// assert_eq!(iter.peek_nth(1), None);
70    /// ```
71    pub fn peek_nth(&mut self, n: usize) -> Option<&I::Item> {
72        let unbuffered_items = (n + 1).saturating_sub(self.buf.len());
73
74        self.buf.extend(self.iter.by_ref().take(unbuffered_items));
75
76        self.buf.get(n)
77    }
78
79    /// Returns a mutable reference to the `nth` value without advancing the iterator.
80    ///
81    /// # Examples
82    ///
83    /// Basic usage:
84    ///
85    /// ```
86    /// use itertools::peek_nth;
87    ///
88    /// let xs = vec![1, 2, 3, 4, 5];
89    /// let mut iter = peek_nth(xs.into_iter());
90    ///
91    /// assert_eq!(iter.peek_nth_mut(0), Some(&mut 1));
92    /// assert_eq!(iter.next(), Some(1));
93    ///
94    /// // The iterator does not advance even if we call `peek_nth_mut` multiple times
95    /// assert_eq!(iter.peek_nth_mut(0), Some(&mut 2));
96    /// assert_eq!(iter.peek_nth_mut(1), Some(&mut 3));
97    /// assert_eq!(iter.next(), Some(2));
98    ///
99    /// // Peek into the iterator and set the value behind the mutable reference.
100    /// if let Some(p) = iter.peek_nth_mut(1) {
101    ///     assert_eq!(*p, 4);
102    ///     *p = 9;
103    /// }
104    ///
105    /// // The value we put in reappears as the iterator continues.
106    /// assert_eq!(iter.next(), Some(3));
107    /// assert_eq!(iter.next(), Some(9));
108    ///
109    /// // Calling `peek_nth_mut` past the end of the iterator will return `None`
110    /// assert_eq!(iter.peek_nth_mut(1), None);
111    /// ```
112    pub fn peek_nth_mut(&mut self, n: usize) -> Option<&mut I::Item> {
113        let unbuffered_items = (n + 1).saturating_sub(self.buf.len());
114
115        self.buf.extend(self.iter.by_ref().take(unbuffered_items));
116
117        self.buf.get_mut(n)
118    }
119
120    /// Works exactly like the `next_if` method in [`std::iter::Peekable`].
121    pub fn next_if(&mut self, func: impl FnOnce(&I::Item) -> bool) -> Option<I::Item> {
122        match self.next() {
123            Some(item) if func(&item) => Some(item),
124            Some(item) => {
125                self.buf.push_front(item);
126                None
127            }
128            _ => None,
129        }
130    }
131
132    /// Works exactly like the `next_if_eq` method in [`std::iter::Peekable`].
133    pub fn next_if_eq<T>(&mut self, expected: &T) -> Option<I::Item>
134    where
135        T: ?Sized,
136        I::Item: PartialEq<T>,
137    {
138        self.next_if(|next| next == expected)
139    }
140}
141
142impl<I> Iterator for PeekNth<I>
143where
144    I: Iterator,
145{
146    type Item = I::Item;
147
148    fn next(&mut self) -> Option<Self::Item> {
149        self.buf.pop_front().or_else(|| self.iter.next())
150    }
151
152    fn size_hint(&self) -> (usize, Option<usize>) {
153        size_hint::add_scalar(self.iter.size_hint(), self.buf.len())
154    }
155
156    fn fold<B, F>(self, mut init: B, mut f: F) -> B
157    where
158        F: FnMut(B, Self::Item) -> B,
159    {
160        init = self.buf.into_iter().fold(init, &mut f);
161        self.iter.fold(init, f)
162    }
163}
164
165impl<I> ExactSizeIterator for PeekNth<I> where I: ExactSizeIterator {}
166
167impl<I> PeekingNext for PeekNth<I>
168where
169    I: Iterator,
170{
171    fn peeking_next<F>(&mut self, accept: F) -> Option<Self::Item>
172    where
173        F: FnOnce(&Self::Item) -> bool,
174    {
175        self.peek().filter(|item| accept(item))?;
176        self.next()
177    }
178}
itertools/peek_nth.rs

itertools/
peek_nth.rs
