Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 1 | use alloc::vec::Vec; |
| 2 | |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 3 | use crate::size_hint; |
| 4 | |
| 5 | /// An iterator adaptor that allows putting multiple |
| 6 | /// items in front of the iterator. |
| 7 | /// |
| 8 | /// Iterator element type is `I::Item`. |
| 9 | #[derive(Debug, Clone)] |
| 10 | pub struct PutBackN<I: Iterator> { |
| 11 | top: Vec<I::Item>, |
| 12 | iter: I, |
| 13 | } |
| 14 | |
| 15 | /// Create an iterator where you can put back multiple values to the front |
| 16 | /// of the iteration. |
| 17 | /// |
| 18 | /// Iterator element type is `I::Item`. |
| 19 | pub fn put_back_n<I>(iterable: I) -> PutBackN<I::IntoIter> |
| 20 | where I: IntoIterator |
| 21 | { |
| 22 | PutBackN { |
| 23 | top: Vec::new(), |
| 24 | iter: iterable.into_iter(), |
| 25 | } |
| 26 | } |
| 27 | |
| 28 | impl<I: Iterator> PutBackN<I> { |
| 29 | /// Puts x in front of the iterator. |
| 30 | /// The values are yielded in order of the most recently put back |
| 31 | /// values first. |
| 32 | /// |
| 33 | /// ```rust |
| 34 | /// use itertools::put_back_n; |
| 35 | /// |
| 36 | /// let mut it = put_back_n(1..5); |
| 37 | /// it.next(); |
| 38 | /// it.put_back(1); |
| 39 | /// it.put_back(0); |
| 40 | /// |
| 41 | /// assert!(itertools::equal(it, 0..5)); |
| 42 | /// ``` |
| 43 | #[inline] |
| 44 | pub fn put_back(&mut self, x: I::Item) { |
| 45 | self.top.push(x); |
| 46 | } |
| 47 | } |
| 48 | |
| 49 | impl<I: Iterator> Iterator for PutBackN<I> { |
| 50 | type Item = I::Item; |
| 51 | #[inline] |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 52 | fn next(&mut self) -> Option<Self::Item> { |
| 53 | self.top.pop().or_else(|| self.iter.next()) |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 54 | } |
| 55 | |
| 56 | #[inline] |
| 57 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 58 | size_hint::add_scalar(self.iter.size_hint(), self.top.len()) |
| 59 | } |
| 60 | } |
| 61 | |