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