From 5b776556feae34dc2a675f99996a169f57f7e289 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Thu, 26 Sep 2024 17:46:58 +0200
Subject: [PATCH] Add ParallelIteratorExt

---
 milli/src/update/new/items_pool.rs | 32 ++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/milli/src/update/new/items_pool.rs b/milli/src/update/new/items_pool.rs
index e90ce97db..c57bc86f1 100644
--- a/milli/src/update/new/items_pool.rs
+++ b/milli/src/update/new/items_pool.rs
@@ -1,4 +1,36 @@
 use crossbeam_channel::{Receiver, Sender, TryRecvError};
+use rayon::iter::{MapInit, ParallelIterator};
+
+pub trait ParallelIteratorExt: ParallelIterator {
+    fn try_map_try_init<F, INIT, T, E, R>(
+        self,
+        init: INIT,
+        map_op: F,
+    ) -> MapInit<
+        Self,
+        impl Fn() -> Result<T, Option<E>> + Sync + Send + Clone,
+        impl Fn(&mut Result<T, Option<E>>, Self::Item) -> Result<R, Option<E>> + Sync + Send + Clone,
+    >
+    where
+        E: Send,
+        F: Fn(&mut T, Self::Item) -> Result<R, E> + Sync + Send + Clone,
+        INIT: Fn() -> Result<T, E> + Sync + Send + Clone,
+        R: Send,
+    {
+        self.map_init(
+            move || match init() {
+                Ok(t) => Ok(t),
+                Err(err) => Err(Some(err)),
+            },
+            move |maybe_t, item| match maybe_t {
+                Ok(t) => map_op(t, item).map_err(Some),
+                Err(maybe_err) => Err(maybe_err.take()),
+            },
+        )
+    }
+}
+
+impl<T: ParallelIterator> ParallelIteratorExt for T {}
 
 /// A pool of items that can be pull and generated on demand.
 pub struct ItemsPool<F, T, E>