Mercurial > hg
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