feat fork_join: add map_array and map_list

src/fork_join.ml

@@ -175,6 +175,38 @@ let all_init ?chunk_size n f : _ list =
       | None -> assert false
       | Some x -> x)
 
+let map_array ?chunk_size f arr : _ array =
+  let n = Array.length arr in
+  let res = Array.make n None in
+
+  for_ ?chunk_size n (fun low high ->
+      for i = low to high do
+        res.(i) <- Some (f arr.(i))
+      done);
+
+  (* get all results *)
+  Array.map
+    (function
+      | None -> assert false
+      | Some x -> x)
+    res
+
+let map_list ?chunk_size f (l : _ list) : _ list =
+  let arr = Array.of_list l in
+  let n = Array.length arr in
+  let res = Array.make n None in
+
+  for_ ?chunk_size n (fun low high ->
+      for i = low to high do
+        res.(i) <- Some (f arr.(i))
+      done);
+
+  (* get all results *)
+  List.init n (fun i ->
+      match res.(i) with
+      | None -> assert false
+      | Some x -> x)
+
 type 'a commutative_monoid = {
   neutral: unit -> 'a;  (** Neutral element *)
   combine: 'a -> 'a -> 'a;  (** Combine two items. *)

src/fork_join.mli

@@ -96,6 +96,16 @@ val all_init : ?chunk_size:int -> int -> (int -> 'a) -> 'a list
     @since 0.3
     {b NOTE} this is only available on OCaml 5. *)
 
+val map_array : ?chunk_size:int -> ('a -> 'b) -> 'a array -> 'b array
+(** [map_array f arr] is like [Array.map f arr], but runs in parallel.
+    @since NEXT_RELEASE
+    {b NOTE} this is only available on OCaml 5. *)
+
+val map_list : ?chunk_size:int -> ('a -> 'b) -> 'a list -> 'b list
+(** [map_list f l] is like [List.map f l], but runs in parallel.
+    @since NEXT_RELEASE
+    {b NOTE} this is only available on OCaml 5. *)
+
 type 'a commutative_monoid = {
   neutral: unit -> 'a;  (** Neutral element *)
   combine: 'a -> 'a -> 'a;  (** Combine two items. *)