| 1 | use super::size_hint; |
| 2 | |
| 3 | use alloc::collections::VecDeque; |
| 4 | use alloc::rc::Rc; |
| 5 | use std::cell::RefCell; |
| 6 | |
| 7 | /// Common buffer object for the two tee halves |
| 8 | #[derive (Debug)] |
| 9 | struct 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)] |
| 21 | pub struct Tee<I> |
| 22 | where |
| 23 | I: Iterator, |
| 24 | { |
| 25 | rcbuffer: Rc<RefCell<TeeBuffer<I::Item, I>>>, |
| 26 | id: bool, |
| 27 | } |
| 28 | |
| 29 | pub fn new<I>(iter: I) -> (Tee<I>, Tee<I>) |
| 30 | where |
| 31 | I: Iterator, |
| 32 | { |
| 33 | let buffer: TeeBuffer<::Item, …> = TeeBuffer { |
| 34 | backlog: VecDeque::new(), |
| 35 | iter, |
| 36 | owner: false, |
| 37 | }; |
| 38 | let t1: Tee = Tee { |
| 39 | rcbuffer: Rc::new(RefCell::new(buffer)), |
| 40 | id: true, |
| 41 | }; |
| 42 | let t2: Tee = Tee { |
| 43 | rcbuffer: t1.rcbuffer.clone(), |
| 44 | id: false, |
| 45 | }; |
| 46 | (t1, t2) |
| 47 | } |
| 48 | |
| 49 | impl<I> Iterator for Tee<I> |
| 50 | where |
| 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 | |
| 88 | impl<I> ExactSizeIterator for Tee<I> |
| 89 | where |
| 90 | I: ExactSizeIterator, |
| 91 | I::Item: Clone, |
| 92 | { |
| 93 | } |
| 94 | |