itertools/
tee.rs

1use super::size_hint;
2
3use std::cell::RefCell;
4use std::collections::VecDeque;
5use std::rc::Rc;
6
7/// Common buffer object for the two tee halves
8#[derive(Debug)]
9struct TeeBuffer<A, I> {
10    backlog: VecDeque<A>,
11    iter: I,
12    /// The owner field indicates which id should read from the backlog
13    owner: bool,
14}
15
16/// One half of an iterator pair where both return the same elements.
17///
18/// See [`.tee()`](../trait.Itertools.html#method.tee) for more information.
19#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
20#[derive(Debug)]
21pub struct Tee<I>
22    where I: Iterator
23{
24    rcbuffer: Rc<RefCell<TeeBuffer<I::Item, I>>>,
25    id: bool,
26}
27
28pub fn new<I>(iter: I) -> (Tee<I>, Tee<I>)
29    where I: Iterator
30{
31    let buffer = TeeBuffer{backlog: VecDeque::new(), iter: iter, owner: false};
32    let t1 = Tee{rcbuffer: Rc::new(RefCell::new(buffer)), id: true};
33    let t2 = Tee{rcbuffer: t1.rcbuffer.clone(), id: false};
34    (t1, t2)
35}
36
37impl<I> Iterator for Tee<I>
38    where I: Iterator,
39          I::Item: Clone
40{
41    type Item = I::Item;
42    fn next(&mut self) -> Option<I::Item> {
43        // .borrow_mut may fail here -- but only if the user has tied some kind of weird
44        // knot where the iterator refers back to itself.
45        let mut buffer = self.rcbuffer.borrow_mut();
46        if buffer.owner == self.id {
47            match buffer.backlog.pop_front() {
48                None => {}
49                some_elt => return some_elt,
50            }
51        }
52        match buffer.iter.next() {
53            None => None,
54            Some(elt) => {
55                buffer.backlog.push_back(elt.clone());
56                buffer.owner = !self.id;
57                Some(elt)
58            }
59        }
60    }
61
62    fn size_hint(&self) -> (usize, Option<usize>) {
63        let buffer = self.rcbuffer.borrow();
64        let sh = buffer.iter.size_hint();
65
66        if buffer.owner == self.id {
67            let log_len = buffer.backlog.len();
68            size_hint::add_scalar(sh, log_len)
69        } else {
70            sh
71        }
72    }
73}
74
75impl<I> ExactSizeIterator for Tee<I>
76    where I: ExactSizeIterator,
77          I::Item: Clone
78{}