itertools/
tee.rs

1use super::size_hint;
2
3use alloc::collections::VecDeque;
4use alloc::rc::Rc;
5use std::cell::RefCell;
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()`](crate::Itertools::tee) for more information.
19#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
20#[derive(Debug)]
21pub struct Tee<I>
22where
23    I: Iterator,
24{
25    rcbuffer: Rc<RefCell<TeeBuffer<I::Item, I>>>,
26    id: bool,
27}
28
29pub fn new<I>(iter: I) -> (Tee<I>, Tee<I>)
30where
31    I: Iterator,
32{
33    let buffer = TeeBuffer {
34        backlog: VecDeque::new(),
35        iter,
36        owner: false,
37    };
38    let t1 = Tee {
39        rcbuffer: Rc::new(RefCell::new(buffer)),
40        id: true,
41    };
42    let t2 = Tee {
43        rcbuffer: t1.rcbuffer.clone(),
44        id: false,
45    };
46    (t1, t2)
47}
48
49impl<I> Iterator for Tee<I>
50where
51    I: Iterator,
52    I::Item: Clone,
53{
54    type Item = I::Item;
55    fn next(&mut self) -> Option<Self::Item> {
56        // .borrow_mut may fail here -- but only if the user has tied some kind of weird
57        // knot where the iterator refers back to itself.
58        let mut buffer = self.rcbuffer.borrow_mut();
59        if buffer.owner == self.id {
60            match buffer.backlog.pop_front() {
61                None => {}
62                some_elt => return some_elt,
63            }
64        }
65        match buffer.iter.next() {
66            None => None,
67            Some(elt) => {
68                buffer.backlog.push_back(elt.clone());
69                buffer.owner = !self.id;
70                Some(elt)
71            }
72        }
73    }
74
75    fn size_hint(&self) -> (usize, Option<usize>) {
76        let buffer = self.rcbuffer.borrow();
77        let sh = buffer.iter.size_hint();
78
79        if buffer.owner == self.id {
80            let log_len = buffer.backlog.len();
81            size_hint::add_scalar(sh, log_len)
82        } else {
83            sh
84        }
85    }
86}
87
88impl<I> ExactSizeIterator for Tee<I>
89where
90    I: ExactSizeIterator,
91    I::Item: Clone,
92{
93}
itertools/tee.rs

itertools/
tee.rs
