Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 1 | use std::iter::{Fuse, FusedIterator}; |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 2 | use super::size_hint; |
| 3 | |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 4 | pub trait IntersperseElement<Item> { |
| 5 | fn generate(&mut self) -> Item; |
| 6 | } |
| 7 | |
| 8 | #[derive(Debug, Clone)] |
| 9 | pub struct IntersperseElementSimple<Item>(Item); |
| 10 | |
| 11 | impl<Item: Clone> IntersperseElement<Item> for IntersperseElementSimple<Item> { |
| 12 | fn generate(&mut self) -> Item { |
| 13 | self.0.clone() |
| 14 | } |
| 15 | } |
| 16 | |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 17 | /// An iterator adaptor to insert a particular value |
| 18 | /// between each element of the adapted iterator. |
| 19 | /// |
| 20 | /// Iterator element type is `I::Item` |
| 21 | /// |
| 22 | /// This iterator is *fused*. |
| 23 | /// |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 24 | /// See [`.intersperse()`](crate::Itertools::intersperse) for more information. |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 25 | pub type Intersperse<I> = IntersperseWith<I, IntersperseElementSimple<<I as Iterator>::Item>>; |
| 26 | |
| 27 | /// Create a new Intersperse iterator |
| 28 | pub fn intersperse<I>(iter: I, elt: I::Item) -> Intersperse<I> |
| 29 | where I: Iterator, |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 30 | { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 31 | intersperse_with(iter, IntersperseElementSimple(elt)) |
| 32 | } |
| 33 | |
| 34 | impl<Item, F: FnMut()->Item> IntersperseElement<Item> for F { |
| 35 | fn generate(&mut self) -> Item { |
| 36 | self() |
| 37 | } |
| 38 | } |
| 39 | |
| 40 | /// An iterator adaptor to insert a particular value created by a function |
| 41 | /// between each element of the adapted iterator. |
| 42 | /// |
| 43 | /// Iterator element type is `I::Item` |
| 44 | /// |
| 45 | /// This iterator is *fused*. |
| 46 | /// |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 47 | /// See [`.intersperse_with()`](crate::Itertools::intersperse_with) for more information. |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 48 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed"] |
| 49 | #[derive(Clone, Debug)] |
| 50 | pub struct IntersperseWith<I, ElemF> |
| 51 | where I: Iterator, |
| 52 | { |
| 53 | element: ElemF, |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 54 | iter: Fuse<I>, |
| 55 | peek: Option<I::Item>, |
| 56 | } |
| 57 | |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 58 | /// Create a new IntersperseWith iterator |
| 59 | pub fn intersperse_with<I, ElemF>(iter: I, elt: ElemF) -> IntersperseWith<I, ElemF> |
| 60 | where I: Iterator, |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 61 | { |
| 62 | let mut iter = iter.fuse(); |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 63 | IntersperseWith { |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 64 | peek: iter.next(), |
| 65 | iter, |
| 66 | element: elt, |
| 67 | } |
| 68 | } |
| 69 | |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 70 | impl<I, ElemF> Iterator for IntersperseWith<I, ElemF> |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 71 | where I: Iterator, |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 72 | ElemF: IntersperseElement<I::Item> |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 73 | { |
| 74 | type Item = I::Item; |
| 75 | #[inline] |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 76 | fn next(&mut self) -> Option<Self::Item> { |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 77 | if self.peek.is_some() { |
| 78 | self.peek.take() |
| 79 | } else { |
| 80 | self.peek = self.iter.next(); |
| 81 | if self.peek.is_some() { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 82 | Some(self.element.generate()) |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 83 | } else { |
| 84 | None |
| 85 | } |
| 86 | } |
| 87 | } |
| 88 | |
| 89 | fn size_hint(&self) -> (usize, Option<usize>) { |
| 90 | // 2 * SH + { 1 or 0 } |
| 91 | let has_peek = self.peek.is_some() as usize; |
| 92 | let sh = self.iter.size_hint(); |
| 93 | size_hint::add_scalar(size_hint::add(sh, sh), has_peek) |
| 94 | } |
| 95 | |
| 96 | fn fold<B, F>(mut self, init: B, mut f: F) -> B where |
| 97 | Self: Sized, F: FnMut(B, Self::Item) -> B, |
| 98 | { |
| 99 | let mut accum = init; |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 100 | |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 101 | if let Some(x) = self.peek.take() { |
| 102 | accum = f(accum, x); |
| 103 | } |
| 104 | |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 105 | let element = &mut self.element; |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 106 | |
| 107 | self.iter.fold(accum, |
| 108 | |accum, x| { |
Joel Galenson | 6f79871 | 2021-04-01 17:03:06 -0700 | [diff] [blame] | 109 | let accum = f(accum, element.generate()); |
Jakub Kotur | a425e55 | 2020-12-21 17:28:15 +0100 | [diff] [blame] | 110 | let accum = f(accum, x); |
| 111 | accum |
| 112 | }) |
| 113 | } |
| 114 | } |
Joel Galenson | b593e25 | 2021-06-21 13:15:57 -0700 | [diff] [blame] | 115 | |
| 116 | impl<I, ElemF> FusedIterator for IntersperseWith<I, ElemF> |
| 117 | where I: Iterator, |
| 118 | ElemF: IntersperseElement<I::Item> |
| 119 | {} |