From db3386815a80bab1953e310c61329dd8497c9026 Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Wed, 27 Nov 2024 14:22:17 +0200
Subject: [PATCH] Hmm

---
 lib/picos_std.structured/run.ml | 53 ++++++++++++++++++++++-----------
 test/test_structured.ml         |  8 +++--
 2 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/lib/picos_std.structured/run.ml b/lib/picos_std.structured/run.ml
index d2564582..12352612 100644
--- a/lib/picos_std.structured/run.ml
+++ b/lib/picos_std.structured/run.ml
@@ -51,45 +51,64 @@ type _ tdt =
 
 external lo_as_atomic : [ `Range ] tdt -> int Atomic.t = "%identity"
 
-let rec for_out (Range r as range : [ `Range ] tdt) action =
+let rec for_out (Range r as range : [ `Range ] tdt) per_fiber action =
   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
+    let lo_after = lo_before + ((n + 1) asr 1) in
+    if Atomic.compare_and_set (lo_as_atomic range) lo_before lo_after then begin
+      per_fiber := lo_before;
+      while !per_fiber < lo_after do
+        let i = !per_fiber in
+        per_fiber := i + 1;
+        action i
+      done;
+      for_out range per_fiber action
     end
     else begin
       (* Contention, bail out... *)
       match r.parent with
       | Empty -> ()
-      | Range _ as range -> for_out range action
+      | Range _ as range -> for_out range per_fiber action
     end
   end
   else
-    match r.parent with Empty -> () | Range _ as range -> for_out range action
+    match r.parent with
+    | Empty -> ()
+    | Range _ as range -> for_out range per_fiber action
 
-let rec for_in bundle (Range r as range : [ `Range ] tdt) action =
+let rec for_in bundle (Range r as range : [ `Range ] tdt) per_fiber action =
   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
-  end
+  if n <= 2 then for_out range per_fiber action
   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);
+      Bundle.fork bundle (fun () -> for_in_enter bundle range action);
       let child = Range { _lo = lo_before; hi = lo_after; parent = range } in
-      for_in bundle child action
+      for_in bundle child per_fiber action
     end
-    else for_in bundle range action
+    else for_in bundle range per_fiber action
+
+and for_in_enter bundle (Range r as range : [ `Range ] tdt) action =
+  let per_fiber = ref r.hi in
+  let effc (type a) :
+      a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option = function
+    | Fiber.Spawn _ -> None
+    | _ ->
+        if !per_fiber < r.hi then begin
+          let range = Range { _lo = !per_fiber; hi = r.hi; parent = Empty } in
+          per_fiber := r.hi;
+          Bundle.fork bundle (fun () -> for_in_enter bundle range action)
+        end;
+        None
+  in
+  let handler = Effect.Deep.{ effc } in
+  Effect.Deep.try_with (for_in bundle range per_fiber) action handler
 
 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
+      Bundle.join_after @@ fun bundle -> for_in_enter bundle range action
diff --git a/test/test_structured.ml b/test/test_structured.ml
index 5e871e4d..63395c5b 100644
--- a/test/test_structured.ml
+++ b/test/test_structured.ml
@@ -263,9 +263,13 @@ let test_for_n_basic () =
   @@ 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)
+      Bytes.set bytes i (Char.chr 0)
+    done;
+    Run.for_n n (fun i ->
+        Bytes.set bytes i (Char.chr (Char.code (Bytes.get bytes i) + 1)));
+    for i = 0 to n - 1 do
+      assert (Bytes.get bytes i = Char.chr 1)
     done
   done