Fiber: add map-reduce helpers

src/dune_engine/action_builder.ml

@@ -118,9 +118,16 @@ let rec eval : type a m. a t -> m eval_mode -> (a * m) Memo.t =
     res, Deps_or_facts.union_all mode deps
   | All_unit ts ->
     let open Memo.O in
-    let+ res = Memo.parallel_map ts ~f:(fun t -> eval t mode) in
-    let deps = List.map res ~f:snd in
-    (), Deps_or_facts.union_all mode deps
+    let+ deps =
+      Memo.map_reduce
+        ts
+        ~f:(fun t ->
+          let+ _, deps = eval t mode in
+          deps)
+        ~empty:(Deps_or_facts.empty mode)
+        ~combine:(Deps_or_facts.union mode)
+    in
+    (), deps
   | Of_memo memo ->
     let open Memo.O in
     let+ x = memo in

src/dune_engine/reflection.ml

@@ -24,34 +24,39 @@ end = struct
         "expand-alias"
         ~input:(module Alias)
         (fun alias ->
-           let* l =
+           let* deps =
              Load_rules.get_alias_definition alias
-             >>= Memo.parallel_map ~f:(fun (loc, definition) ->
-               Memo.push_stack_frame
-                 (fun () ->
-                    Action_builder.evaluate_and_collect_deps
-                      (Build_system.dep_on_alias_definition definition)
-                    >>| snd)
-                 ~human_readable_description:(fun () -> Alias.describe alias ~loc))
+             >>= Memo.map_reduce
+                   ~empty:Dep.Set.empty
+                   ~combine:Dep.Set.union
+                   ~f:(fun (loc, definition) ->
+                     Memo.push_stack_frame
+                       (fun () ->
+                          Action_builder.evaluate_and_collect_deps
+                            (Build_system.dep_on_alias_definition definition)
+                          >>| snd)
+                       ~human_readable_description:(fun () -> Alias.describe alias ~loc))
            in
-           let deps = List.fold_left l ~init:Dep.Set.empty ~f:Dep.Set.union in
            Expand.deps deps)
     in
     Memo.exec memo
   ;;
 
   let deps deps =
-    Memo.parallel_map (Dep.Set.to_list deps) ~f:(fun (dep : Dep.t) ->
-      match dep with
-      | File p -> Memo.return (Path.Set.singleton p)
-      | File_selector g ->
-        let+ filenames = Build_system.eval_pred g in
-        (* Alas, we can't use filename sets here because we end up putting paths coming
+    Memo.map_reduce
+      (Dep.Set.to_list deps)
+      ~empty:Path.Set.empty
+      ~combine:Path.Set.union
+      ~f:(fun (dep : Dep.t) ->
+        match dep with
+        | File p -> Memo.return (Path.Set.singleton p)
+        | File_selector g ->
+          let+ filenames = Build_system.eval_pred g in
+          (* Alas, we can't use filename sets here because we end up putting paths coming
            from different directories together. *)
-        Path.Set.of_list (Filename_set.to_list filenames)
-      | Alias a -> Expand.alias a
-      | Env _ | Universe -> Memo.return Path.Set.empty)
-    >>| Path.Set.union_all
+          Path.Set.of_list (Filename_set.to_list filenames)
+        | Alias a -> Expand.alias a
+        | Env _ | Universe -> Memo.return Path.Set.empty)
   ;;
 end
 

src/dune_rules/coq/coq_rules.ml

@@ -397,12 +397,16 @@ let directories_of_lib ~sctx lib =
 
 let setup_native_theory_includes ~sctx ~theories_deps ~theory_dirs =
   Resolve.Memo.bind theories_deps ~f:(fun theories_deps ->
-    let+ l =
-      Memo.parallel_map theories_deps ~f:(fun lib ->
-        let+ theory_dirs = directories_of_lib ~sctx lib in
-        Path.Build.Set.of_list theory_dirs)
+    let+ theory_dirs =
+      Memo.map_reduce
+        theories_deps
+        ~empty:theory_dirs
+        ~combine:Path.Build.Set.union
+        ~f:(fun lib ->
+          let+ theory_dirs = directories_of_lib ~sctx lib in
+          Path.Build.Set.of_list theory_dirs)
     in
-    Resolve.return (Path.Build.Set.union_all (theory_dirs :: l)))
+    Resolve.return theory_dirs)
 ;;
 
 let coqc_native_flags ~sctx ~dir ~theories_deps ~theory_dirs ~(mode : Coq_mode.t) =
@@ -992,15 +996,19 @@ let setup_coqdoc_rules ~sctx ~dir ~theories_deps (s : Coq_stanza.Theory.t) coq_m
              @@
              let open Memo.O in
              let+ deps =
-               Memo.parallel_map theories_deps ~f:(fun theory ->
-                 let+ theory_dirs = directories_of_lib ~sctx theory in
-                 Dep.Set.of_list_map theory_dirs ~f:(fun dir ->
-                   (* TODO *)
-                   Glob.of_string_exn Loc.none "*.glob"
-                   |> File_selector.of_glob ~dir:(Path.build dir)
-                   |> Dep.file_selector))
+               Memo.map_reduce
+                 theories_deps
+                 ~empty:Dep.Set.empty
+                 ~combine:Dep.Set.union
+                 ~f:(fun theory ->
+                   let+ theory_dirs = directories_of_lib ~sctx theory in
+                   Dep.Set.of_list_map theory_dirs ~f:(fun dir ->
+                     (* TODO *)
+                     Glob.of_string_exn Loc.none "*.glob"
+                     |> File_selector.of_glob ~dir:(Path.build dir)
+                     |> Dep.file_selector))
              in
-             Command.Args.Hidden_deps (Dep.Set.union_all deps)
+             Command.Args.Hidden_deps deps
            in
            let mode_flag =
              match mode with

src/dune_rules/lib.ml

@@ -2634,9 +2634,11 @@ module DB = struct
     let open Memo.O in
     let* l =
       Memo.Lazy.force t.all
-      >>= Memo.parallel_map ~f:(find t)
-      >>| List.filter_opt
-      >>| Set.of_list
+      >>= Memo.map_reduce ~empty:Set.empty ~combine:Set.union ~f:(fun name ->
+        find t name
+        >>| function
+        | None -> Set.empty
+        | Some lib -> Set.singleton lib)
     in
     match recursive, t.parent with
     | true, Some t ->

src/dune_rules/lock_rules.ml

@@ -439,19 +439,21 @@ let scan_lock_directory =
         ]
     | Ok entries ->
       Fs_memo.Dir_contents.to_list entries
-      |> Memo.parallel_map ~f:(fun (entry, kind) ->
-        let path = Path.Outside_build_dir.relative dir entry in
-        match (kind : File_kind.t) with
-        | S_REG -> Memo.return (Path.Set.singleton (Path.outside_build_dir path))
-        | S_DIR -> scan path
-        | kind ->
-          User_error.raise
-            [ Pp.textf
-                "Lock directory contains file %S with unsupported kind %S"
-                (Path.Outside_build_dir.to_string_maybe_quoted path)
-                (File_kind.to_string kind)
-            ])
-      >>| Path.Set.union_all
+      |> Memo.map_reduce
+           ~empty:Path.Set.empty
+           ~combine:Path.Set.union
+           ~f:(fun (entry, kind) ->
+             let path = Path.Outside_build_dir.relative dir entry in
+             match (kind : File_kind.t) with
+             | S_REG -> Memo.return (Path.Set.singleton (Path.outside_build_dir path))
+             | S_DIR -> scan path
+             | kind ->
+               User_error.raise
+                 [ Pp.textf
+                     "Lock directory contains file %S with unsupported kind %S"
+                     (Path.Outside_build_dir.to_string_maybe_quoted path)
+                     (File_kind.to_string kind)
+                 ])
   in
   fun lock_dir_path ->
     let+ files = scan lock_dir_path in

src/dune_rules/mdx.ml

@@ -135,8 +135,14 @@ module Deps = struct
     | Error e -> Memo.return (Error e)
     | Ok (dirs, files) ->
       let dep_set = Dep.Set.of_files files in
-      let+ l = Memo.parallel_map dirs ~f:(fun dir -> Source_deps.files dir >>| fst) in
-      Ok (Dep.Set.union_all (dep_set :: l))
+      let+ deps =
+        Memo.map_reduce
+          dirs
+          ~f:(fun dir -> Source_deps.files dir >>| fst)
+          ~empty:dep_set
+          ~combine:Dep.Set.union
+      in
+      Ok deps
   ;;
 end
 

src/dune_rules/merlin/merlin.ml

@@ -697,18 +697,18 @@ module Unprocessed = struct
   ;;
 
   let add_lib_dirs sctx ~for_ libs =
-    Memo.parallel_map libs ~f:(fun lib ->
-      let+ dirs = src_dirs sctx lib ~for_ in
-      lib, dirs)
-    >>| List.fold_left
-          ~init:(Path.Set.empty, Path.Set.empty)
-          ~f:(fun (src_dirs, obj_dirs) (lib, more_src_dirs) ->
-            ( Path.Set.union src_dirs more_src_dirs
-            , let public_cmi_dir =
-                let info = Lib.info lib in
-                obj_dir_of_lib `Public for_ (Lib_info.obj_dir info)
-              in
-              Path.Set.add obj_dirs public_cmi_dir ))
+    Memo.map_reduce
+      libs
+      ~empty:(Path.Set.empty, Path.Set.empty)
+      ~combine:(fun (src_dirs1, obj_dirs1) (src_dirs2, obj_dirs2) ->
+        Path.Set.union src_dirs1 src_dirs2, Path.Set.union obj_dirs1 obj_dirs2)
+      ~f:(fun lib ->
+        let+ src_dirs = src_dirs sctx lib ~for_ in
+        let obj_dir =
+          let info = Lib.info lib in
+          obj_dir_of_lib `Public for_ (Lib_info.obj_dir info)
+        in
+        src_dirs, Path.Set.singleton obj_dir)
     |> Action_builder.of_memo
   ;;
 

src/dune_rules/pkg_rules.ml

@@ -508,10 +508,11 @@ module Pkg = struct
             | _ -> if skip_file name then Skip else Left relative)
         in
         let acc = Path.Local.Set.of_list files |> Path.Local.Set.union acc in
-        let+ dirs =
-          Memo.parallel_map dirs ~f:(fun dir -> loop root Path.Local.Set.empty dir)
-        in
-        Path.Local.Set.union_all (acc :: dirs)
+        Memo.map_reduce
+          dirs
+          ~f:(fun dir -> loop root Path.Local.Set.empty dir)
+          ~empty:acc
+          ~combine:Path.Local.Set.union
     in
     (match t.info.source with
      | None -> Memo.return None

src/dune_rules/rocq/rocq_rules.ml

@@ -336,12 +336,16 @@ let directories_of_lib ~sctx lib =
 
 let setup_native_theory_includes ~sctx ~theories_deps ~theory_dirs =
   Resolve.Memo.bind theories_deps ~f:(fun theories_deps ->
-    let+ l =
-      Memo.parallel_map theories_deps ~f:(fun lib ->
-        let+ theory_dirs = directories_of_lib ~sctx lib in
-        Path.Build.Set.of_list theory_dirs)
+    let+ theory_dirs =
+      Memo.map_reduce
+        theories_deps
+        ~empty:theory_dirs
+        ~combine:Path.Build.Set.union
+        ~f:(fun lib ->
+          let+ theory_dirs = directories_of_lib ~sctx lib in
+          Path.Build.Set.of_list theory_dirs)
     in
-    Resolve.return (Path.Build.Set.union_all (theory_dirs :: l)))
+    Resolve.return theory_dirs)
 ;;
 
 let rocqc_native_flags ~sctx ~dir ~theories_deps ~theory_dirs ~(mode : Rocq_mode.t) =
@@ -945,15 +949,19 @@ let setup_rocqdoc_rules ~sctx ~dir ~theories_deps (s : Rocq_stanza.Theory.t) roc
              @@
              let open Memo.O in
              let+ deps =
-               Memo.parallel_map theories_deps ~f:(fun theory ->
-                 let+ theory_dirs = directories_of_lib ~sctx theory in
-                 Dep.Set.of_list_map theory_dirs ~f:(fun dir ->
-                   (* TODO *)
-                   Glob.of_string_exn Loc.none "*.glob"
-                   |> File_selector.of_glob ~dir:(Path.build dir)
-                   |> Dep.file_selector))
+               Memo.map_reduce
+                 theories_deps
+                 ~empty:Dep.Set.empty
+                 ~combine:Dep.Set.union
+                 ~f:(fun theory ->
+                   let+ theory_dirs = directories_of_lib ~sctx theory in
+                   Dep.Set.of_list_map theory_dirs ~f:(fun dir ->
+                     (* TODO *)
+                     Glob.of_string_exn Loc.none "*.glob"
+                     |> File_selector.of_glob ~dir:(Path.build dir)
+                     |> Dep.file_selector))
              in
-             Command.Args.Hidden_deps (Dep.Set.union_all deps)
+             Command.Args.Hidden_deps deps
            in
            let mode_flag =
              match mode with

src/fiber/src/cache.ml

@@ -36,14 +36,16 @@ let find_or_add t key ~f =
 let to_table t =
   let* () = return () in
   let+ values_by_key =
-    Table.foldi t.table ~init:[] ~f:(fun key ivar acc ->
-      (let+ value_result = Ivar.read ivar in
-       match value_result with
-       | Ok value -> Some (key, value)
-       | Error _ -> None)
-      :: acc)
-    |> all_concurrently
-    >>| List.filter_opt
+    Table.foldi t.table ~init:[] ~f:(fun key ivar acc -> (key, ivar) :: acc)
+    |> map_reduce
+         ~empty:Appendable_list.empty
+         ~combine:Appendable_list.( @ )
+         ~f:(fun (key, ivar) ->
+           let+ value_result = Ivar.read ivar in
+           match value_result with
+           | Ok value -> Appendable_list.singleton (key, value)
+           | Error _ -> Appendable_list.empty)
+    >>| Appendable_list.to_list
   in
   let table = Table.create t.key_module 1 in
   List.iter values_by_key ~f:(fun (key, value) -> Table.add_exn table key value);

src/fiber/src/core.ml

@@ -133,6 +133,15 @@ let rec nfork_seq left_over x (seq : _ Seq.t) f =
      | eff -> Fork (eff, fun () -> nfork_seq left_over y seq f))
 ;;
 
+let rec nfork_array a i f =
+  if i = Array.length a - 1
+  then f a.(i)
+  else (
+    match apply f a.(i) with
+    | End_of_fiber () -> nfork_array a (i + 1) f
+    | eff -> Fork (eff, fun () -> nfork_array a (i + 1) f))
+;;
+
 let parallel_iter_seq (seq : _ Seq.t) ~f k =
   match seq () with
   | Nil -> k ()
@@ -146,6 +155,51 @@ let parallel_iter_seq (seq : _ Seq.t) ~f k =
     nfork_seq left_over x seq f
 ;;
 
+let map_reduce_seq (seq : _ Seq.t) ~f ~empty ~combine k =
+  match seq () with
+  | Nil -> k empty
+  | Cons (x, seq) ->
+    let current = ref empty in
+    let running = ref 1 in
+    let f x =
+      f x (fun y ->
+        current := combine !current y;
+        decr running;
+        if !running = 0 then k !current else end_of_fiber)
+    in
+    nfork_seq running x seq f
+;;
+
+let map_reduce_array a ~f ~empty ~combine k =
+  match Array.length a with
+  | 0 -> k empty
+  | len ->
+    let current = ref empty in
+    let running = ref len in
+    let f x =
+      f x (fun y ->
+        current := combine !current y;
+        decr running;
+        if !running = 0 then k !current else end_of_fiber)
+    in
+    nfork_array a 0 f
+;;
+
+let map_reduce l ~f ~empty ~combine k =
+  match l with
+  | [] -> k empty
+  | x :: l ->
+    let current = ref empty in
+    let running = ref (List.length l + 1) in
+    let f x =
+      f x (fun y ->
+        current := combine !current y;
+        decr running;
+        if !running = 0 then k !current else end_of_fiber)
+    in
+    nfork x l f
+;;
+
 type ('a, 'b) fork_and_join_state =
   | Nothing_yet
   | Got_a of 'a