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.
with_position<I>(iter: I) -> WithPosition<I> where I: 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 
next(&mut self) -> Option<Self::Item>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 
size_hint(&self) -> (usize, Option<usize>)90     fn size_hint(&self) -> (usize, Option<usize>) {
91         self.peekable.size_hint()
92     }
93 
fold<B, F>(mut self, mut init: B, mut f: F) -> B where F: FnMut(B, Self::Item) -> B,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