changeset 51154:aba622c7dc7e

rust: add a utility function to merge ordered fallible iterators Adding a function merge_join_results_by, a version of itertools::merge_join_by that works on "fallible" iterators (iterators that can produce errors)
author Arseniy Alekseyev <aalekseyev@janestreet.com>
date Wed, 15 Nov 2023 18:41:33 +0000
parents 976403c95ba3
children ac3859a8b796
files rust/hg-core/src/utils.rs
diffstat 1 files changed, 40 insertions(+), 0 deletions(-) [+]
line wrap: on
line diff
--- a/rust/hg-core/src/utils.rs	Mon Oct 16 18:56:40 2023 +0100
+++ b/rust/hg-core/src/utils.rs	Wed Nov 15 18:41:33 2023 +0000
@@ -11,7 +11,10 @@
 use crate::utils::hg_path::HgPath;
 use im_rc::ordmap::DiffItem;
 use im_rc::ordmap::OrdMap;
+use itertools::EitherOrBoth;
+use itertools::Itertools;
 use std::cell::Cell;
+use std::cmp::Ordering;
 use std::fmt;
 use std::{io::Write, ops::Deref};
 
@@ -499,6 +502,43 @@
     })
 }
 
+/// Like `itertools::merge_join_by`, but merges fallible iterators.
+///
+/// The callback is only used for Ok values. Errors are passed through as-is.
+/// Errors compare less than Ok values, which makes the error handling
+/// conservative.
+pub fn merge_join_results_by<'a, I1, I2, F, A, B, E>(
+    iter1: I1,
+    iter2: I2,
+    f: F,
+) -> impl Iterator<Item = Result<EitherOrBoth<A, B>, E>> + 'a
+where
+    I1: Iterator<Item = Result<A, E>> + 'a,
+    I2: Iterator<Item = Result<B, E>> + 'a,
+    F: FnMut(&A, &B) -> Ordering + 'a,
+{
+    let mut g = f;
+    iter1
+        .merge_join_by(iter2, move |i1, i2| match i1 {
+            Err(_) => Ordering::Less,
+            Ok(i1) => match i2 {
+                Err(_) => Ordering::Greater,
+                Ok(i2) => g(i1, i2),
+            },
+        })
+        .map(|result| match result {
+            EitherOrBoth::Left(Err(e)) => Err(e),
+            EitherOrBoth::Right(Err(e)) => Err(e),
+            EitherOrBoth::Both(Err(e), _) => Err(e),
+            EitherOrBoth::Both(_, Err(e)) => Err(e),
+            EitherOrBoth::Left(Ok(v)) => Ok(EitherOrBoth::Left(v)),
+            EitherOrBoth::Right(Ok(v)) => Ok(EitherOrBoth::Right(v)),
+            EitherOrBoth::Both(Ok(v1), Ok(v2)) => {
+                Ok(EitherOrBoth::Both(v1, v2))
+            }
+        })
+}
+
 /// Force the global rayon threadpool to not exceed 16 concurrent threads
 /// unless the user has specified a value.
 /// This is a stop-gap measure until we figure out why using more than 16