Skip to main content

itertools/
with_position.rs

1use std::fmt;
2use 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"]
10pub struct WithPosition<I>
11where
12    I: Iterator,
13{
14    handled_first: bool,
15    peekable: Peekable<Fuse<I>>,
16}
17
18impl<I> fmt::Debug for WithPosition<I>
19where
20    I: Iterator,
21    Peekable<Fuse<I>>: fmt::Debug,
22{
23    debug_fmt_fields!(WithPosition, handled_first, peekable);
24}
25
26impl<I> Clone for WithPosition<I>
27where
28    I: Clone + Iterator,
29    I::Item: Clone,
30{
31    clone_fields!(handled_first, peekable);
32}
33
34/// Create a new `WithPosition` iterator.
35pub fn with_position<I>(iter: I) -> WithPosition<I>
36where
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)]
50pub struct Position {
51    /// This is the initial element (also true if there's exactly one element)
52    pub is_first: bool,
53    /// This is the final element (also true if there's exactly one element)
54    pub is_last: bool,
55}
56
57impl Position {
58    /// This is the first and the last element at the same time, and there are no more elements
59    pub fn is_exactly_one(self) -> bool {
60        self.is_first && self.is_last
61    }
62
63    /// This is neither first nor last element, and there will be more elements
64    pub fn is_middle(self) -> bool {
65        !self.is_first && !self.is_last
66    }
67
68    /// This is the initial element (also true if there's exactly one element)
69    pub fn is_first(self) -> bool {
70        self.is_first
71    }
72
73    /// This is the final element (also true if there's exactly one element)
74    pub fn is_last(self) -> bool {
75        self.is_last
76    }
77}
78
79impl<I: Iterator> Iterator for WithPosition<I> {
80    type Item = (Position, I::Item);
81
82    fn next(&mut self) -> Option<Self::Item> {
83        let item = self.peekable.next()?;
84
85        let is_last = self.peekable.peek().is_none();
86        let is_first = !self.handled_first;
87        self.handled_first = true;
88
89        Some((Position { is_first, is_last }, item))
90    }
91
92    fn size_hint(&self) -> (usize, Option<usize>) {
93        self.peekable.size_hint()
94    }
95
96    fn fold<B, F>(mut self, mut init: B, mut f: F) -> B
97    where
98        F: FnMut(B, Self::Item) -> B,
99    {
100        if let Some(mut head) = self.peekable.next() {
101            if !self.handled_first {
102                // The current head is `First` or `Only`,
103                // it depends if there is another item or not.
104                match self.peekable.next() {
105                    Some(second) => {
106                        let first = std::mem::replace(&mut head, second);
107                        let position = Position {
108                            is_first: true,
109                            is_last: false,
110                        };
111                        init = f(init, (position, first));
112                    }
113                    None => {
114                        let position = Position {
115                            is_first: true,
116                            is_last: true,
117                        };
118                        return f(init, (position, head));
119                    }
120                }
121            }
122            // Have seen the first item, and there's something left.
123            init = self.peekable.fold(init, |acc, mut item| {
124                std::mem::swap(&mut head, &mut item);
125                let position = Position {
126                    is_first: false,
127                    is_last: false,
128                };
129                f(acc, (position, item))
130            });
131            let position = Position {
132                is_first: false,
133                is_last: true,
134            };
135            // The "head" is now the last item.
136            init = f(init, (position, head));
137        }
138        init
139    }
140}
141
142impl<I> ExactSizeIterator for WithPosition<I> where I: ExactSizeIterator {}
143
144impl<I: Iterator> FusedIterator for WithPosition<I> {}