summaryrefslogtreecommitdiffstats
path: root/third_party/rust/itertools/src/kmerge_impl.rs
diff options
context:
space:
mode:
authorDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 00:47:55 +0000
committerDaniel Baumann <daniel.baumann@progress-linux.org>2024-04-19 00:47:55 +0000
commit26a029d407be480d791972afb5975cf62c9360a6 (patch)
treef435a8308119effd964b339f76abb83a57c29483 /third_party/rust/itertools/src/kmerge_impl.rs
parentInitial commit. (diff)
downloadfirefox-26a029d407be480d791972afb5975cf62c9360a6.tar.xz
firefox-26a029d407be480d791972afb5975cf62c9360a6.zip
Adding upstream version 124.0.1.upstream/124.0.1
Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
Diffstat (limited to 'third_party/rust/itertools/src/kmerge_impl.rs')
-rw-r--r--third_party/rust/itertools/src/kmerge_impl.rs227
1 files changed, 227 insertions, 0 deletions
diff --git a/third_party/rust/itertools/src/kmerge_impl.rs b/third_party/rust/itertools/src/kmerge_impl.rs
new file mode 100644
index 0000000000..509d5fc6a3
--- /dev/null
+++ b/third_party/rust/itertools/src/kmerge_impl.rs
@@ -0,0 +1,227 @@
+use crate::size_hint;
+use crate::Itertools;
+
+use alloc::vec::Vec;
+use std::iter::FusedIterator;
+use std::mem::replace;
+use std::fmt;
+
+/// Head element and Tail iterator pair
+///
+/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
+/// first items (which are guaranteed to exist).
+///
+/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
+/// `KMerge` into a min-heap.
+#[derive(Debug)]
+struct HeadTail<I>
+ where I: Iterator
+{
+ head: I::Item,
+ tail: I,
+}
+
+impl<I> HeadTail<I>
+ where I: Iterator
+{
+ /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
+ fn new(mut it: I) -> Option<HeadTail<I>> {
+ let head = it.next();
+ head.map(|h| {
+ HeadTail {
+ head: h,
+ tail: it,
+ }
+ })
+ }
+
+ /// Get the next element and update `head`, returning the old head in `Some`.
+ ///
+ /// Returns `None` when the tail is exhausted (only `head` then remains).
+ fn next(&mut self) -> Option<I::Item> {
+ if let Some(next) = self.tail.next() {
+ Some(replace(&mut self.head, next))
+ } else {
+ None
+ }
+ }
+
+ /// Hints at the size of the sequence, same as the `Iterator` method.
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ size_hint::add_scalar(self.tail.size_hint(), 1)
+ }
+}
+
+impl<I> Clone for HeadTail<I>
+ where I: Iterator + Clone,
+ I::Item: Clone
+{
+ clone_fields!(head, tail);
+}
+
+/// Make `data` a heap (min-heap w.r.t the sorting).
+fn heapify<T, S>(data: &mut [T], mut less_than: S)
+ where S: FnMut(&T, &T) -> bool
+{
+ for i in (0..data.len() / 2).rev() {
+ sift_down(data, i, &mut less_than);
+ }
+}
+
+/// Sift down element at `index` (`heap` is a min-heap wrt the ordering)
+fn sift_down<T, S>(heap: &mut [T], index: usize, mut less_than: S)
+ where S: FnMut(&T, &T) -> bool
+{
+ debug_assert!(index <= heap.len());
+ let mut pos = index;
+ let mut child = 2 * pos + 1;
+ // Require the right child to be present
+ // This allows to find the index of the smallest child without a branch
+ // that wouldn't be predicted if present
+ while child + 1 < heap.len() {
+ // pick the smaller of the two children
+ // use arithmetic to avoid an unpredictable branch
+ child += less_than(&heap[child+1], &heap[child]) as usize;
+
+ // sift down is done if we are already in order
+ if !less_than(&heap[child], &heap[pos]) {
+ return;
+ }
+ heap.swap(pos, child);
+ pos = child;
+ child = 2 * pos + 1;
+ }
+ // Check if the last (left) child was an only child
+ // if it is then it has to be compared with the parent
+ if child + 1 == heap.len() && less_than(&heap[child], &heap[pos]) {
+ heap.swap(pos, child);
+ }
+}
+
+/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
+/// If all base iterators are sorted (ascending), the result is sorted.
+///
+/// Iterator element type is `I::Item`.
+///
+/// See [`.kmerge()`](crate::Itertools::kmerge) for more information.
+pub type KMerge<I> = KMergeBy<I, KMergeByLt>;
+
+pub trait KMergePredicate<T> {
+ fn kmerge_pred(&mut self, a: &T, b: &T) -> bool;
+}
+
+#[derive(Clone, Debug)]
+pub struct KMergeByLt;
+
+impl<T: PartialOrd> KMergePredicate<T> for KMergeByLt {
+ fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
+ a < b
+ }
+}
+
+impl<T, F: FnMut(&T, &T)->bool> KMergePredicate<T> for F {
+ fn kmerge_pred(&mut self, a: &T, b: &T) -> bool {
+ self(a, b)
+ }
+}
+
+/// Create an iterator that merges elements of the contained iterators using
+/// the ordering function.
+///
+/// [`IntoIterator`] enabled version of [`Itertools::kmerge`].
+///
+/// ```
+/// use itertools::kmerge;
+///
+/// for elt in kmerge(vec![vec![0, 2, 4], vec![1, 3, 5], vec![6, 7]]) {
+/// /* loop body */
+/// }
+/// ```
+pub fn kmerge<I>(iterable: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
+ where I: IntoIterator,
+ I::Item: IntoIterator,
+ <<I as IntoIterator>::Item as IntoIterator>::Item: PartialOrd
+{
+ kmerge_by(iterable, KMergeByLt)
+}
+
+/// An iterator adaptor that merges an abitrary number of base iterators
+/// according to an ordering function.
+///
+/// Iterator element type is `I::Item`.
+///
+/// See [`.kmerge_by()`](crate::Itertools::kmerge_by) for more
+/// information.
+#[must_use = "iterator adaptors are lazy and do nothing unless consumed"]
+pub struct KMergeBy<I, F>
+ where I: Iterator,
+{
+ heap: Vec<HeadTail<I>>,
+ less_than: F,
+}
+
+impl<I, F> fmt::Debug for KMergeBy<I, F>
+ where I: Iterator + fmt::Debug,
+ I::Item: fmt::Debug,
+{
+ debug_fmt_fields!(KMergeBy, heap);
+}
+
+/// Create an iterator that merges elements of the contained iterators.
+///
+/// [`IntoIterator`] enabled version of [`Itertools::kmerge_by`].
+pub fn kmerge_by<I, F>(iterable: I, mut less_than: F)
+ -> KMergeBy<<I::Item as IntoIterator>::IntoIter, F>
+ where I: IntoIterator,
+ I::Item: IntoIterator,
+ F: KMergePredicate<<<I as IntoIterator>::Item as IntoIterator>::Item>,
+{
+ let iter = iterable.into_iter();
+ let (lower, _) = iter.size_hint();
+ let mut heap: Vec<_> = Vec::with_capacity(lower);
+ heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
+ heapify(&mut heap, |a, b| less_than.kmerge_pred(&a.head, &b.head));
+ KMergeBy { heap, less_than }
+}
+
+impl<I, F> Clone for KMergeBy<I, F>
+ where I: Iterator + Clone,
+ I::Item: Clone,
+ F: Clone,
+{
+ clone_fields!(heap, less_than);
+}
+
+impl<I, F> Iterator for KMergeBy<I, F>
+ where I: Iterator,
+ F: KMergePredicate<I::Item>
+{
+ type Item = I::Item;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ if self.heap.is_empty() {
+ return None;
+ }
+ let result = if let Some(next) = self.heap[0].next() {
+ next
+ } else {
+ self.heap.swap_remove(0).head
+ };
+ let less_than = &mut self.less_than;
+ sift_down(&mut self.heap, 0, |a, b| less_than.kmerge_pred(&a.head, &b.head));
+ Some(result)
+ }
+
+ fn size_hint(&self) -> (usize, Option<usize>) {
+ #[allow(deprecated)] //TODO: once msrv hits 1.51. replace `fold1` with `reduce`
+ self.heap.iter()
+ .map(|i| i.size_hint())
+ .fold1(size_hint::add)
+ .unwrap_or((0, Some(0)))
+ }
+}
+
+impl<I, F> FusedIterator for KMergeBy<I, F>
+ where I: Iterator,
+ F: KMergePredicate<I::Item>
+{}