Add more structured Run operations

polytypic · polytypic · commit d4e02107c6b8 · 2024-11-27T02:17:12.000+02:00
diff --git a/lib/picos_std.structured/picos_std_structured.mli b/lib/picos_std.structured/picos_std_structured.mli
@@ -398,6 +398,13 @@ module Run : sig
           with Control.Terminate -> ()
       ]}
       but treats the list of actions as a single computation. *)
+
+  val for_n : int -> (int -> unit) -> unit
+  (** [for_n n action], when [0 < n], calls [action i] for each integer [i] from
+      [0] to [n-1] using up to [n] fibers, including the current fiber.  The
+      actual number of fibers used can be much less than [n] in case the calls
+      do not block, the calls return quickly, and/or the scheduler doesn't
+      provide parallelism. *)
 end
 
 (** {1 Examples}
diff --git a/lib/picos_std.structured/run.ml b/lib/picos_std.structured/run.ml
@@ -37,3 +37,59 @@ let run actions wrap =
 
 let all actions = run actions wrap_all
 let any actions = run actions wrap_any
+
+(* *)
+
+type _ tdt =
+  | Empty : [> `Empty ] tdt
+  | Range : {
+      mutable _lo : int;
+      hi : int;
+      parent : [ `Empty | `Range ] tdt;
+    }
+      -> [> `Range ] tdt
+
+external lo_as_atomic : [ `Range ] tdt -> int Atomic.t = "%identity"
+
+let rec for_out (Range r as range : [ `Range ] tdt) action backoff =
+  let lo_before = Atomic.get (lo_as_atomic range) in
+  let n = r.hi - lo_before in
+  if 0 < n then begin
+    if Atomic.compare_and_set (lo_as_atomic range) lo_before (lo_before + 1)
+    then begin
+      action lo_before;
+      for_out range action Backoff.default
+    end
+    else for_out range action (Backoff.once backoff)
+  end
+  else
+    match r.parent with
+    | Empty -> ()
+    | Range _ as range -> for_out range action Backoff.default
+
+let rec for_in bundle (Range r as range : [ `Range ] tdt) action backoff =
+  let lo_before = Atomic.get (lo_as_atomic range) in
+  let n = r.hi - lo_before in
+  if n <= 1 then begin
+    if n = 1 && Atomic.compare_and_set (lo_as_atomic range) lo_before r.hi then
+      action lo_before;
+    match r.parent with
+    | Empty -> ()
+    | Range _ as range -> for_out range action Backoff.default
+  end
+  else
+    let lo_after = lo_before + (n asr 1) in
+    if Atomic.compare_and_set (lo_as_atomic range) lo_before lo_after then begin
+      Bundle.fork bundle (fun () -> for_in bundle range action Backoff.default);
+      let child = Range { _lo = lo_before; hi = lo_after; parent = range } in
+      for_in bundle child action Backoff.default
+    end
+    else for_in bundle range action (Backoff.once backoff)
+
+let for_n n action =
+  if 0 < n then
+    if n = 1 then action 0
+    else
+      let range = Range { _lo = 0; hi = n; parent = Empty } in
+      Bundle.join_after @@ fun bundle ->
+      for_in bundle range action Backoff.default
diff --git a/test/dune b/test/dune
@@ -249,6 +249,7 @@
  (libraries
   alcotest
   picos
+  picos.domain
   picos_aux.mpscq
   picos_std.finally
   picos_std.structured
diff --git a/test/test_structured.ml b/test/test_structured.ml
@@ -258,6 +258,17 @@ let test_race_any () =
   (* This is non-deterministic and may need to changed if flaky *)
   assert (Atomic.get winner = 1)
 
+let test_for_n_basic () =
+  Test_scheduler.run ~max_domains:(Picos_domain.recommended_domain_count ())
+  @@ fun () ->
+  for n = 0 to 128 do
+    let bytes = Bytes.create n in
+    Run.for_n n (fun i -> Bytes.set bytes i (Char.chr i));
+    for i = 0 to n - 1 do
+      assert (Bytes.get bytes i = Char.chr i)
+    done
+  done
+
 let () =
   [
     ( "Bundle",
@@ -280,9 +291,13 @@ let () =
           test_error_in_promise_terminates;
         Alcotest.test_case "can wait promises" `Quick test_can_wait_promises;
         Alcotest.test_case "can select promises" `Quick test_can_select_promises;
+      ] );
+    ( "Run",
+      [
         Alcotest.test_case "any and all errors" `Quick test_any_and_all_errors;
         Alcotest.test_case "any and all returns" `Quick test_any_and_all_returns;
         Alcotest.test_case "race any" `Quick test_race_any;
+        Alcotest.test_case "for_n basic" `Quick test_for_n_basic;
       ] );
   ]
   |> Alcotest.run "Picos_structured"
