blob: a0d79b03632dbc7fab0cd7c69afaeba3d54079fc [file] [log] [blame]
Jakub Kotura425e552020-12-21 17:28:15 +01001use std::iter::Fuse;
2use super::size_hint;
3
Joel Galenson6f798712021-04-01 17:03:06 -07004pub trait IntersperseElement<Item> {
5 fn generate(&mut self) -> Item;
6}
7
8#[derive(Debug, Clone)]
9pub struct IntersperseElementSimple<Item>(Item);
10
11impl<Item: Clone> IntersperseElement<Item> for IntersperseElementSimple<Item> {
12 fn generate(&mut self) -> Item {
13 self.0.clone()
14 }
15}
16
Jakub Kotura425e552020-12-21 17:28:15 +010017/// 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///
24/// See [`.intersperse()`](../trait.Itertools.html#method.intersperse) for more information.
Joel Galenson6f798712021-04-01 17:03:06 -070025pub type Intersperse<I> = IntersperseWith<I, IntersperseElementSimple<<I as Iterator>::Item>>;
26
27/// Create a new Intersperse iterator
28pub fn intersperse<I>(iter: I, elt: I::Item) -> Intersperse<I>
29 where I: Iterator,
Jakub Kotura425e552020-12-21 17:28:15 +010030{
Joel Galenson6f798712021-04-01 17:03:06 -070031 intersperse_with(iter, IntersperseElementSimple(elt))
32}
33
34impl<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///
47/// See [`.intersperse_with()`](../trait.Itertools.html#method.intersperse_with) for more information.
48#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
49#[derive(Clone, Debug)]
50pub struct IntersperseWith<I, ElemF>
51 where I: Iterator,
52{
53 element: ElemF,
Jakub Kotura425e552020-12-21 17:28:15 +010054 iter: Fuse<I>,
55 peek: Option<I::Item>,
56}
57
Joel Galenson6f798712021-04-01 17:03:06 -070058/// Create a new IntersperseWith iterator
59pub fn intersperse_with<I, ElemF>(iter: I, elt: ElemF) -> IntersperseWith<I, ElemF>
60 where I: Iterator,
Jakub Kotura425e552020-12-21 17:28:15 +010061{
62 let mut iter = iter.fuse();
Joel Galenson6f798712021-04-01 17:03:06 -070063 IntersperseWith {
Jakub Kotura425e552020-12-21 17:28:15 +010064 peek: iter.next(),
65 iter,
66 element: elt,
67 }
68}
69
Joel Galenson6f798712021-04-01 17:03:06 -070070impl<I, ElemF> Iterator for IntersperseWith<I, ElemF>
Jakub Kotura425e552020-12-21 17:28:15 +010071 where I: Iterator,
Joel Galenson6f798712021-04-01 17:03:06 -070072 ElemF: IntersperseElement<I::Item>
Jakub Kotura425e552020-12-21 17:28:15 +010073{
74 type Item = I::Item;
75 #[inline]
Joel Galenson6f798712021-04-01 17:03:06 -070076 fn next(&mut self) -> Option<Self::Item> {
Jakub Kotura425e552020-12-21 17:28:15 +010077 if self.peek.is_some() {
78 self.peek.take()
79 } else {
80 self.peek = self.iter.next();
81 if self.peek.is_some() {
Joel Galenson6f798712021-04-01 17:03:06 -070082 Some(self.element.generate())
Jakub Kotura425e552020-12-21 17:28:15 +010083 } 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 Galenson6f798712021-04-01 17:03:06 -0700100
Jakub Kotura425e552020-12-21 17:28:15 +0100101 if let Some(x) = self.peek.take() {
102 accum = f(accum, x);
103 }
104
Joel Galenson6f798712021-04-01 17:03:06 -0700105 let element = &mut self.element;
Jakub Kotura425e552020-12-21 17:28:15 +0100106
107 self.iter.fold(accum,
108 |accum, x| {
Joel Galenson6f798712021-04-01 17:03:06 -0700109 let accum = f(accum, element.generate());
Jakub Kotura425e552020-12-21 17:28:15 +0100110 let accum = f(accum, x);
111 accum
112 })
113 }
114}