library/alloc/src/collections/btree/merge_iter.rs

   1 use core::cmp::Ordering;
   2 use core::fmt::{self, Debug};
   3 use core::iter::FusedIterator;
   4
   5 /// Core of an iterator that merges the output of two strictly ascending iterators,
   6 /// for instance a union or a symmetric difference.
   7 pub struct MergeIterInner<I: Iterator> {
   8     a: I,
   9     b: I,
  10     peeked: Option<Peeked<I>>,
  11 }
  12
  13 /// Benchmarks faster than wrapping both iterators in a Peekable,
  14 /// probably because we can afford to impose a FusedIterator bound.
  15 #[derive(Clone, Debug)]
  16 enum Peeked<I: Iterator> {
  17     A(I::Item),
  18     B(I::Item),
  19 }
  20
  21 impl<I: Iterator> Clone for MergeIterInner<I>
  22 where
  23     I: Clone,
  24     I::Item: Clone,
  25 {
  26     fn clone(&self) -> Self {
  27         Self { a: self.a.clone(), b: self.b.clone(), peeked: self.peeked.clone() }
  28     }
  29 }
  30
  31 impl<I: Iterator> Debug for MergeIterInner<I>
  32 where
  33     I: Debug,
  34     I::Item: Debug,
  35 {
  36     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
  37         f.debug_tuple("MergeIterInner").field(&self.a).field(&self.b).field(&self.peeked).finish()
  38     }
  39 }
  40
  41 impl<I: Iterator> MergeIterInner<I> {
  42     /// Creates a new core for an iterator merging a pair of sources.
  43     pub fn new(a: I, b: I) -> Self {
  44         MergeIterInner { a, b, peeked: None }
  45     }
  46
  47     /// Returns the next pair of items stemming from the pair of sources
  48     /// being merged. If both returned options contain a value, that value
  49     /// is equal and occurs in both sources. If one of the returned options
  50     /// contains a value, that value doesn't occur in the other source (or
  51     /// the sources are not strictly ascending). If neither returned option
  52     /// contains a value, iteration has finished and subsequent calls will
  53     /// return the same empty pair.
  54     pub fn nexts<Cmp: Fn(&I::Item, &I::Item) -> Ordering>(
  55         &mut self,
  56         cmp: Cmp,
  57     ) -> (Option<I::Item>, Option<I::Item>)
  58     where
  59         I: FusedIterator,
  60     {
  61         let mut a_next;
  62         let mut b_next;
  63         match self.peeked.take() {
  64             Some(Peeked::A(next)) => {
  65                 a_next = Some(next);
  66                 b_next = self.b.next();
  67             }
  68             Some(Peeked::B(next)) => {
  69                 b_next = Some(next);
  70                 a_next = self.a.next();
  71             }
  72             None => {
  73                 a_next = self.a.next();
  74                 b_next = self.b.next();
  75             }
  76         }
  77         if let (Some(ref a1), Some(ref b1)) = (&a_next, &b_next) {
  78             match cmp(a1, b1) {
  79                 Ordering::Less => self.peeked = b_next.take().map(Peeked::B),
  80                 Ordering::Greater => self.peeked = a_next.take().map(Peeked::A),
  81                 Ordering::Equal => (),
  82             }
  83         }
  84         (a_next, b_next)
  85     }
  86
  87     /// Returns a pair of upper bounds for the `size_hint` of the final iterator.
  88     pub fn lens(&self) -> (usize, usize)
  89     where
  90         I: ExactSizeIterator,
  91     {
  92         match self.peeked {
  93             Some(Peeked::A(_)) => (1 + self.a.len(), self.b.len()),
  94             Some(Peeked::B(_)) => (self.a.len(), 1 + self.b.len()),
  95             _ => (self.a.len(), self.b.len()),
  96         }
  97     }
  98 }