vendor/itertools/src/tee.rs

   1 use super::size_hint;
   2
   3 use std::cell::RefCell;
   4 use std::collections::VecDeque;
   5 use std::rc::Rc;
   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()`](../trait.Itertools.html#method.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 I: Iterator
  23 {
  24     rcbuffer: Rc<RefCell<TeeBuffer<I::Item, I>>>,
  25     id: bool,
  26 }
  27
  28 pub 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
  37 impl<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
  75 impl<I> ExactSizeIterator for Tee<I>
  76     where I: ExactSizeIterator,
  77           I::Item: Clone
  78 {}