1 | use std::fmt; |
2 | use std::iter::{Fuse, FusedIterator, Peekable}; |
3 | |
4 | /// An iterator adaptor that wraps each element in an [`Position`]. |
5 | /// |
6 | /// Iterator element type is `(Position, I::Item)`. |
7 | /// |
8 | /// See [`.with_position()`](crate::Itertools::with_position) for more information. |
9 | #[must_use = "iterator adaptors are lazy and do nothing unless consumed" ] |
10 | pub struct WithPosition<I> |
11 | where |
12 | I: Iterator, |
13 | { |
14 | handled_first: bool, |
15 | peekable: Peekable<Fuse<I>>, |
16 | } |
17 | |
18 | impl<I> fmt::Debug for WithPosition<I> |
19 | where |
20 | I: Iterator, |
21 | Peekable<Fuse<I>>: fmt::Debug, |
22 | { |
23 | debug_fmt_fields!(WithPosition, handled_first, peekable); |
24 | } |
25 | |
26 | impl<I> Clone for WithPosition<I> |
27 | where |
28 | I: Clone + Iterator, |
29 | I::Item: Clone, |
30 | { |
31 | clone_fields!(handled_first, peekable); |
32 | } |
33 | |
34 | /// Create a new `WithPosition` iterator. |
35 | pub fn with_position<I>(iter: I) -> WithPosition<I> |
36 | where |
37 | I: Iterator, |
38 | { |
39 | WithPosition { |
40 | handled_first: false, |
41 | peekable: iter.fuse().peekable(), |
42 | } |
43 | } |
44 | |
45 | /// The first component of the value yielded by `WithPosition`. |
46 | /// Indicates the position of this element in the iterator results. |
47 | /// |
48 | /// See [`.with_position()`](crate::Itertools::with_position) for more information. |
49 | #[derive (Copy, Clone, Debug, PartialEq, Eq)] |
50 | pub enum Position { |
51 | /// This is the first element. |
52 | First, |
53 | /// This is neither the first nor the last element. |
54 | Middle, |
55 | /// This is the last element. |
56 | Last, |
57 | /// This is the only element. |
58 | Only, |
59 | } |
60 | |
61 | impl<I: Iterator> Iterator for WithPosition<I> { |
62 | type Item = (Position, I::Item); |
63 | |
64 | fn next(&mut self) -> Option<Self::Item> { |
65 | match self.peekable.next() { |
66 | Some(item) => { |
67 | if !self.handled_first { |
68 | // Haven't seen the first item yet, and there is one to give. |
69 | self.handled_first = true; |
70 | // Peek to see if this is also the last item, |
71 | // in which case tag it as `Only`. |
72 | match self.peekable.peek() { |
73 | Some(_) => Some((Position::First, item)), |
74 | None => Some((Position::Only, item)), |
75 | } |
76 | } else { |
77 | // Have seen the first item, and there's something left. |
78 | // Peek to see if this is the last item. |
79 | match self.peekable.peek() { |
80 | Some(_) => Some((Position::Middle, item)), |
81 | None => Some((Position::Last, item)), |
82 | } |
83 | } |
84 | } |
85 | // Iterator is finished. |
86 | None => None, |
87 | } |
88 | } |
89 | |
90 | fn size_hint(&self) -> (usize, Option<usize>) { |
91 | self.peekable.size_hint() |
92 | } |
93 | |
94 | fn fold<B, F>(mut self, mut init: B, mut f: F) -> B |
95 | where |
96 | F: FnMut(B, Self::Item) -> B, |
97 | { |
98 | if let Some(mut head) = self.peekable.next() { |
99 | if !self.handled_first { |
100 | // The current head is `First` or `Only`, |
101 | // it depends if there is another item or not. |
102 | match self.peekable.next() { |
103 | Some(second) => { |
104 | let first = std::mem::replace(&mut head, second); |
105 | init = f(init, (Position::First, first)); |
106 | } |
107 | None => return f(init, (Position::Only, head)), |
108 | } |
109 | } |
110 | // Have seen the first item, and there's something left. |
111 | init = self.peekable.fold(init, |acc, mut item| { |
112 | std::mem::swap(&mut head, &mut item); |
113 | f(acc, (Position::Middle, item)) |
114 | }); |
115 | // The "head" is now the last item. |
116 | init = f(init, (Position::Last, head)); |
117 | } |
118 | init |
119 | } |
120 | } |
121 | |
122 | impl<I> ExactSizeIterator for WithPosition<I> where I: ExactSizeIterator {} |
123 | |
124 | impl<I: Iterator> FusedIterator for WithPosition<I> {} |
125 | |