From 3754d39441644acdd714bd699ec222fa0dfac499 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 | 51 ++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 13 deletions(-)

diff --git a/lib/picos_std.structured/run.ml b/lib/picos_std.structured/run.ml
index d2564582d..ed1c83e5b 100644
--- a/lib/picos_std.structured/run.ml
+++ b/lib/picos_std.structured/run.ml
@@ -51,45 +51,70 @@ 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 + 10 in
+    if Atomic.compare_and_set (lo_as_atomic range) lo_before lo_after then begin
+      for i = lo_before to Int.min lo_after r.hi do
+        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
+    match r.parent with
+    | Empty -> ()
+    | Range _ as range -> for_out range per_fiber action
   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);
+      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 =
+  if r.hi - r._lo = 1 then action r._lo
+  else
+    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