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