vendor/itertools-0.9.0/src/kmerge_impl.rs

   1 use crate::size_hint;
   2 use crate::Itertools;
   3
   4 use std::mem::replace;
   5 use std::fmt;
   6
   7 /// Head element and Tail iterator pair
   8 ///
   9 /// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
  10 /// first items (which are guaranteed to exist).
  11 ///
  12 /// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
  13 /// `KMerge` into a min-heap.
  14 #[derive(Debug)]
  15 struct HeadTail<I>
  16     where I: Iterator
  17 {
  18     head: I::Item,
  19     tail: I,
  20 }
  21
  22 impl<I> HeadTail<I>
  23     where I: Iterator
  24 {
  25     /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
  26     fn new(mut it: I) -> Option<HeadTail<I>> {
  27         let head = it.next();
  28         head.map(|h| {
  29             HeadTail {
  30                 head: h,
  31                 tail: it,
  32             }
  33         })
  34     }
  35
  36     /// Get the next element and update `head`, returning the old head in `Some`.
  37     ///
  38     /// Returns `None` when the tail is exhausted (only `head` then remains).
  39     fn next(&mut self) -> Option<I::Item> {
  40         if let Some(next) = self.tail.next() {
  41             Some(replace(&mut self.head, next))
  42         } else {
  43             None
  44         }
  45     }
  46
  47     /// Hints at the size of the sequence, same as the `Iterator` method.
  48     fn size_hint(&self) -> (usize, Option<usize>) {
  49         size_hint::add_scalar(self.tail.size_hint(), 1)
  50     }
  51 }
  52
  53 impl<I> Clone for HeadTail<I>
  54     where I: Iterator + Clone,
  55           I::Item: Clone
  56 {
  57     clone_fields!(head, tail);
  58 }
  59
  60 /// Make `data` a heap (min-heap w.r.t the sorting).
  61 fn heapify<T, S>(data: &mut [T], mut less_than: S)
  62     where S: FnMut(&T, &T) -> bool
  63 {
  64     for i in (0..data.len() / 2).rev() {
  65         sift_down(data, i, &mut less_than);
  66     }
  67 }
  68
  69 /// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
  70 fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
  71     where S: FnMut(&T, &T) -> bool
  72 {
  73     debug_assert!(index <= heap.len());
  74     let mut pos = index;
  75     let mut child = 2 * pos + 1;
  76     // the `pos` conditional is to avoid a bounds check
  77     while pos < heap.len() && child < heap.len() {
  78         let right = child + 1;
  79
  80         // pick the smaller of the two children
  81         if right < heap.len() && less_than(&heap[right], &heap[child]) {
  82             child = right;
  83         }
  84
  85         // sift down is done if we are already in order
  86         if !less_than(&heap[child], &heap[pos]) {
  87             return;
  88         }
  89         heap.swap(pos, child);
  90         pos = child;
  91         child = 2 * pos + 1;
  92     }
  93 }
  94
  95 /// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
  96 /// If all base iterators are sorted (ascending), the result is sorted.
  97 ///
  98 /// Iterator element type is `I::Item`.
  99 ///
 100 /// See [`.kmerge()`](../trait.Itertools.html#method.kmerge) for more information.
 101 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 102 pub type KMerge<I> = KMergeBy<I, KMergeByLt>;
 103
 104 pub trait KMergePredicate<T> {
 105     fn kmerge_pred(&mut self, a: &T, b: &T) -> bool;
 106 }
 107
 108 #[derive(Clone)]
 109 pub struct KMergeByLt;
 110
 111 impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt {
 112     fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
 113         a < b
 114     }
 115 }
 116
 117 impl<T, F: FnMut(&T, &T)->bool> KMergePredicate<T> for F {
 118     fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
 119         self(a, b)
 120     }
 121 }
 122
 123 /// Create an iterator that merges elements of the contained iterators using
 124 /// the ordering function.
 125 ///
 126 /// Equivalent to `iterable.into_iter().kmerge()`.
 127 ///
 128 /// ```
 129 /// use itertools::kmerge;
 130 ///
 131 /// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) {
 132 ///     /* loop body */
 133 /// }
 134 /// ```
 135 pub fn kmerge<I>(iterable: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
 136     where I: IntoIterator,
 137           I::Item: IntoIterator,
 138           <<I as IntoIterator>::Item as IntoIterator>::Item: PartialOrd
 139 {
 140     kmerge_by(iterable, KMergeByLt)
 141 }
 142
 143 /// An iterator adaptor that merges an abitrary number of base iterators
 144 /// according to an ordering function.
 145 ///
 146 /// Iterator element type is `I::Item`.
 147 ///
 148 /// See [`.kmerge_by()`](../trait.Itertools.html#method.kmerge_by) for more
 149 /// information.
 150 #[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
 151 pub struct KMergeBy<I, F>
 152     where I: Iterator,
 153 {
 154     heap: Vec<HeadTail<I>>,
 155     less_than: F,
 156 }
 157
 158 impl<I, F> fmt::Debug for KMergeBy<I, F>
 159     where I: Iterator + fmt::Debug,
 160           I::Item: fmt::Debug,
 161 {
 162     debug_fmt_fields!(KMergeBy, heap);
 163 }
 164
 165 /// Create an iterator that merges elements of the contained iterators.
 166 ///
 167 /// Equivalent to `iterable.into_iter().kmerge_by(less_than)`.
 168 pub fn kmerge_by<I, F>(iterable: I, mut less_than: F)
 169     -> KMergeBy<<I::Item as IntoIterator>::IntoIter, F>
 170     where I: IntoIterator,
 171           I::Item: IntoIterator,
 172           F: KMergePredicate<<<I as IntoIterator>::Item as IntoIterator>::Item>,
 173 {
 174     let iter = iterable.into_iter();
 175     let (lower, _) = iter.size_hint();
 176     let mut heap: Vec<_> = Vec::with_capacity(lower);
 177     heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
 178     heapify(&mut heap, |a, b| less_than.kmerge_pred(&a.head, &b.head));
 179     KMergeBy { heap, less_than }
 180 }
 181
 182 impl<I, F> Clone for KMergeBy<I, F>
 183     where I: Iterator + Clone,
 184           I::Item: Clone,
 185           F: Clone,
 186 {
 187     clone_fields!(heap, less_than);
 188 }
 189
 190 impl<I, F> Iterator for KMergeBy<I, F>
 191     where I: Iterator,
 192           F: KMergePredicate<I::Item>
 193 {
 194     type Item = I::Item;
 195
 196     fn next(&mut self) -> Option<Self::Item> {
 197         if self.heap.is_empty() {
 198             return None;
 199         }
 200         let result = if let Some(next) = self.heap[0].next() {
 201             next
 202         } else {
 203             self.heap.swap_remove(0).head
 204         };
 205         let less_than = &mut self.less_than;
 206         sift_down(&mut self.heap, 0, |a, b| less_than.kmerge_pred(&a.head, &b.head));
 207         Some(result)
 208     }
 209
 210     fn size_hint(&self) -> (usize, Option<usize>) {
 211         self.heap.iter()
 212                  .map(|i| i.size_hint())
 213                  .fold1(size_hint::add)
 214                  .unwrap_or((0, Some(0)))
 215     }
 216 }