or_iterator/
lib.rs

1//! # Examples
2//!
3//! ```
4//! use or_iterator::OrIterator;
5//!
6//! let v1 = vec![1, 2, 3];
7//! let v2 = vec![4, 5];
8//! let or = v1.iter().or(v2.iter());
9//! assert_eq!(3, or.count());
10//!
11//! let v1 = vec![];
12//! let v2 = vec![4, 5];
13//! let or = v1.iter().or(v2.iter());
14//! assert_eq!(2, or.count());
15//! ```
16
17use core::cmp::max;
18
19pub trait OrIterator: Sized {
20    fn or<I2: Iterator>(self, other: I2) -> Or<Self, I2>;
21}
22
23impl<I1: Iterator> OrIterator for I1 {
24    fn or<I2: Iterator>(self: I1, other: I2) -> Or<I1, I2> {
25        Or {
26            iter1: self,
27            iter2: other,
28            state: Initial,
29        }
30    }
31}
32
33enum State {
34    Initial,
35    InIter1,
36    InIter2,
37}
38
39pub struct Or<I1, I2> {
40    iter1: I1,
41    iter2: I2,
42    state: State,
43}
44
45use State::{InIter1, InIter2, Initial};
46
47impl<I1, I2> Iterator for Or<I1, I2>
48where
49    I1: Iterator,
50    I2: Iterator<Item = I1::Item>,
51{
52    type Item = I1::Item;
53
54    fn next(&mut self) -> Option<Self::Item> {
55        match self.state {
56            Initial => match self.iter1.next() {
57                x @ Some(_) => {
58                    self.state = InIter1;
59                    x
60                }
61                None => {
62                    self.state = InIter2;
63                    self.iter2.next()
64                }
65            },
66            InIter1 => self.iter1.next(),
67            InIter2 => self.iter2.next(),
68        }
69    }
70
71    // implementations below this for performance reasons
72    fn size_hint(&self) -> (usize, Option<usize>) {
73        match self.state {
74            Initial => match self.iter1.size_hint() {
75                (_, Some(0)) => self.iter2.size_hint(),
76                (0, i1_high) => {
77                    let (i2_low, i2_high) = self.iter2.size_hint();
78                    let low = if i2_low > 0 { 1 } else { 0 };
79                    let high = i1_high.and_then(|h1| i2_high.map(|h2| max(h1, h2)));
80                    (low, high)
81                }
82                i1_hint => i1_hint,
83            },
84            InIter1 => self.iter1.size_hint(),
85            InIter2 => self.iter2.size_hint(),
86        }
87    }
88
89    fn fold<B, F: FnMut(B, Self::Item) -> B>(mut self, init: B, mut f: F) -> B {
90        match self.state {
91            Initial => match self.iter1.next() {
92                Some(x) => self.iter1.fold(f(init, x), f),
93                None => self.iter2.fold(init, f),
94            },
95            InIter1 => self.iter1.fold(init, f),
96            InIter2 => self.iter2.fold(init, f),
97        }
98    }
99
100    fn nth(&mut self, n: usize) -> Option<Self::Item> {
101        match self.state {
102            Initial => match self.iter1.next() {
103                x @ Some(_) => {
104                    self.state = InIter1;
105                    if n == 0 {
106                        x
107                    } else {
108                        self.iter1.nth(n - 1)
109                    }
110                }
111                None => {
112                    self.state = InIter1;
113                    self.iter2.nth(n)
114                }
115            },
116            InIter1 => self.iter1.nth(n),
117            InIter2 => self.iter2.nth(n),
118        }
119    }
120
121    fn last(mut self) -> Option<Self::Item> {
122        match self.state {
123            Initial => match self.iter1.next() {
124                x @ Some(_) => self.iter1.last().or(x),
125                None => self.iter2.last(),
126            },
127            InIter1 => self.iter1.last(),
128            InIter2 => self.iter2.last(),
129        }
130    }
131
132    fn count(mut self) -> usize {
133        match self.state {
134            Initial => match self.iter1.next() {
135                Some(_) => 1 + self.iter1.count(),
136                None => self.iter2.count(),
137            },
138            InIter1 => self.iter1.count(),
139            InIter2 => self.iter2.count(),
140        }
141    }
142}
143
144impl<I1, I2> ExactSizeIterator for Or<I1, I2>
145where
146    I1: ExactSizeIterator,
147    I2: ExactSizeIterator<Item = I1::Item>,
148{
149}
150
151#[cfg(test)]
152mod tests {
153    use super::OrIterator;
154    use std::iter::{empty, once};
155
156    #[test]
157    fn basic() {
158        let i1 = once(1);
159        let i2 = once(3);
160        let v = i1.or(i2).collect::<Vec<i32>>();
161        assert_eq!(vec![1], v);
162
163        let i1 = empty();
164        let i2 = once(3);
165        let v = i1.or(i2).collect::<Vec<i32>>();
166        assert_eq!(vec![3], v);
167    }
168
169    #[test]
170    fn nth() {
171        let v1 = vec![1, 2, 3];
172        let v2 = vec![4, 5];
173        let mut or = v1.into_iter().or(v2.into_iter());
174        assert_eq!(Some(1), or.nth(0));
175        assert_eq!(Some(3), or.nth(1));
176        assert_eq!(None, or.nth(3));
177
178        let v1 = vec![];
179        let v2 = vec![4, 5];
180        let mut or = v1.into_iter().or(v2.into_iter());
181        assert_eq!(Some(5), or.nth(1));
182        assert_eq!(None, or.nth(0));
183    }
184
185    #[test]
186    fn len() {
187        let v1 = vec![1, 2, 3];
188        let v2 = vec![4, 5];
189        let or = v1.iter().or(v2.iter());
190        assert_eq!(3, or.len());
191
192        let v1 = vec![];
193        let v2 = vec![4, 5];
194        let or = v1.iter().or(v2.iter());
195        assert_eq!(2, or.len());
196    }
197
198    #[test]
199    fn count() {
200        let v1 = vec![1, 2, 3];
201        let v2 = vec![4, 5];
202        let or = v1.iter().or(v2.iter());
203        assert_eq!(3, or.count());
204
205        let v1 = vec![];
206        let v2 = vec![4, 5];
207        let or = v1.iter().or(v2.iter());
208        assert_eq!(2, or.count());
209    }
210
211    #[test]
212    fn last() {
213        let v1 = vec![1, 2, 3];
214        let v2 = vec![4, 5];
215        let or = v1.into_iter().or(v2.into_iter());
216        assert_eq!(Some(3), or.last());
217
218        let v1 = vec![];
219        let v2 = vec![4, 5];
220        let or = v1.into_iter().or(v2.into_iter());
221        assert_eq!(Some(5), or.last());
222    }
223
224    #[test]
225    fn fold() {
226        let v1 = vec![1, 2, 3];
227        let v2 = vec![4, 5];
228        let or = v1.iter().or(v2.iter());
229        let sum = or.fold(0, |n, v| n + v);
230        assert_eq!(6, sum);
231
232        let v1 = vec![];
233        let v2 = vec![4, 5];
234        let or = v1.iter().or(v2.iter());
235        let sum = or.fold(0, |n, v| n + v);
236        assert_eq!(9, sum);
237    }
238}
or_iterator/lib.rs

or_iterator/
lib.rs
