chain_one/
lib.rs

1use std::iter::{ExactSizeIterator, FusedIterator};
2
3pub struct ChainOne<I>
4where
5    I: Iterator,
6{
7    iter: I,
8    elem: Option<I::Item>,
9}
10
11pub trait WithChainOne: Iterator + Sized {
12    fn chain_one(self, elem: Self::Item) -> ChainOne<Self>;
13}
14
15impl<I> WithChainOne for I
16where
17    I: Iterator,
18{
19    fn chain_one(self, elem: Self::Item) -> ChainOne<Self> {
20        ChainOne {
21            iter: self,
22            elem: Some(elem),
23        }
24    }
25}
26
27impl<I> Iterator for ChainOne<I>
28where
29    I: Iterator,
30{
31    type Item = I::Item;
32
33    #[inline]
34    fn next(&mut self) -> Option<Self::Item> {
35        self.iter.next().or_else(|| self.elem.take())
36    }
37
38    #[inline]
39    fn size_hint(&self) -> (usize, Option<usize>) {
40        let (low, hi) = self.iter.size_hint();
41        let second_len = if self.elem.is_some() { 1 } else { 0 };
42
43        (low + second_len, hi.map(|hi| hi + second_len))
44    }
45}
46
47impl<I> ExactSizeIterator for ChainOne<I> where I: ExactSizeIterator {}
48
49impl<I> FusedIterator for ChainOne<I> where I: FusedIterator {}
50
51impl<I> DoubleEndedIterator for ChainOne<I>
52where
53    I: DoubleEndedIterator,
54{
55    fn next_back(&mut self) -> Option<Self::Item> {
56        self.elem.take().or_else(|| self.iter.next_back())
57    }
58}
59
60#[macro_export]
61macro_rules! iter {
62    () => {
63        ::std::iter::none()
64    };
65    ($first:expr $(, $rest:expr)*$(,)*) => {{
66        let i = ::std::iter::once($first);
67        $(let i = $crate::WithChainOne::chain_one(i, $rest);)*
68        i
69    }};
70}
71
72#[cfg(test)]
73mod tests {
74    #[test]
75    fn iter() {
76        assert_eq!(
77            iter![1, 2, 3, 4, 5].collect::<Vec<_>>(),
78            vec![1, 2, 3, 4, 5]
79        );
80    }
81
82    #[test]
83    fn reversed() {
84        assert_eq!(
85            iter![1, 2, 3, 4, 5].rev().collect::<Vec<_>>(),
86            vec![5, 4, 3, 2, 1]
87        );
88    }
89}