And tweak more for performance

polytypic · polytypic · commit 44f67c58a210 · 2024-10-14T22:36:53.000+03:00
diff --git a/lib/picos_mux.fifo/picos_mux_fifo.ml b/lib/picos_mux.fifo/picos_mux_fifo.ml
@@ -8,7 +8,7 @@ type ready =
   | Resume of
       Fiber.t
       * ((exn * Printexc.raw_backtrace) option, unit) Effect.Deep.continuation
-  | Return of Fiber.Maybe.t * (unit, unit) Effect.Deep.continuation
+  | Return of Fiber.t * (unit, unit) Effect.Deep.continuation
 
 module Mpscq = Picos_aux_mpscq
 
@@ -35,22 +35,21 @@ type t = {
 
 let rec next t =
   match Mpscq.pop_exn t.ready with
-  | Spawn (fiber, main) ->
-      t.fiber <- Fiber.Maybe.of_fiber fiber;
+  | ready -> begin
       t.remaining_quota <- t.quota;
-      Effect.Deep.match_with main fiber t.handler
-  | Return (fiber, k) ->
-      t.fiber <- fiber;
-      t.remaining_quota <- t.quota;
-      Effect.Deep.continue k ()
-  | Continue (fiber, k) ->
-      t.fiber <- Fiber.Maybe.of_fiber fiber;
-      t.remaining_quota <- t.quota;
-      Fiber.continue fiber k ()
-  | Resume (fiber, k) ->
-      t.fiber <- Fiber.Maybe.of_fiber fiber;
-      t.remaining_quota <- t.quota;
-      Fiber.resume fiber k
+      t.fiber <-
+        (match ready with
+        | Spawn (fiber, _)
+        | Continue (fiber, _)
+        | Resume (fiber, _)
+        | Return (fiber, _) ->
+            Fiber.Maybe.of_fiber fiber);
+      match ready with
+      | Spawn (fiber, main) -> Effect.Deep.match_with main fiber t.handler
+      | Return (_, k) -> Effect.Deep.continue k ()
+      | Continue (fiber, k) -> Fiber.continue fiber k ()
+      | Resume (fiber, k) -> Fiber.resume fiber k
+    end
   | exception Mpscq.Empty ->
       t.fiber <- Fiber.Maybe.nothing;
       if t.num_alive_fibers <> 0 then begin
@@ -84,7 +83,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
   let rec t =
     {
       ready;
-      fiber = Fiber.Maybe.nothing;
+      fiber = Fiber.Maybe.of_fiber fiber;
       needs_wakeup;
       mutex;
       condition;
@@ -118,7 +117,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
           Effect.Deep.continue k ()
         end
         else begin
-          Mpscq.push t.ready (Return (t.fiber, k));
+          Mpscq.push t.ready (Return (Fiber.Maybe.to_fiber t.fiber, k));
           next t
         end)
   and discontinue =
@@ -187,8 +186,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
       Condition.broadcast t.condition
     end
   in
-  Mpscq.push t.ready (Spawn (fiber, main));
-  next t
+  Effect.Deep.match_with main fiber t.handler
 
 let run ?quota ?fatal_exn_handler ?(forbid = false) main =
   let computation = Computation.create ~mode:`LIFO () in
diff --git a/lib/picos_mux.multififo/picos_mux_multififo.ml b/lib/picos_mux.multififo/picos_mux_multififo.ml
@@ -12,7 +12,7 @@ type ready =
   | Resume of
       Fiber.t
       * ((exn * Printexc.raw_backtrace) option, unit) Effect.Deep.continuation
-  | Return of Fiber.Maybe.t * (unit, unit) Effect.Deep.continuation
+  | Return of Fiber.t * (unit, unit) Effect.Deep.continuation
 
 type t = {
   mutable num_waiters_non_zero : bool;
@@ -101,19 +101,18 @@ let[@inline] relaxed_wakeup t ~known_not_empty ready =
 
 let exec ready (Per_thread p : per_thread) t =
   p.remaining_quota <- t.quota;
+  p.fiber <-
+    (match ready with
+    | Spawn (fiber, _)
+    | Return (fiber, _)
+    | Continue (fiber, _)
+    | Resume (fiber, _) ->
+        Fiber.Maybe.of_fiber fiber);
   match ready with
-  | Spawn (fiber, main) ->
-      p.fiber <- Fiber.Maybe.of_fiber fiber;
-      Effect.Deep.match_with main fiber t.handler
-  | Return (fiber, k) ->
-      p.fiber <- fiber;
-      Effect.Deep.continue k ()
-  | Continue (fiber, k) ->
-      p.fiber <- Fiber.Maybe.of_fiber fiber;
-      Fiber.continue fiber k ()
-  | Resume (fiber, k) ->
-      p.fiber <- Fiber.Maybe.of_fiber fiber;
-      Fiber.resume fiber k
+  | Spawn (fiber, main) -> Effect.Deep.match_with main fiber t.handler
+  | Return (_, k) -> Effect.Deep.continue k ()
+  | Continue (fiber, k) -> Fiber.continue fiber k ()
+  | Resume (fiber, k) -> Fiber.resume fiber k
 
 let rec next (Per_thread p as pt : per_thread) =
   match Mpmcq.pop_exn p.ready with
@@ -229,7 +228,7 @@ let per_thread context =
           Effect.Deep.continue k ()
         end
         else begin
-          Mpmcq.push p.ready (Return (p.fiber, k));
+          Mpmcq.push p.ready (Return (Fiber.Maybe.to_fiber p.fiber, k));
           next pt
         end)
   and discontinue =
@@ -391,7 +390,7 @@ let runner_on_this_thread t =
 
 let run_fiber ?context:t_opt fiber main =
   let t = match t_opt with None -> context () | Some t -> t in
-  with_per_thread t @@ fun (Per_thread p as pt) ->
+  with_per_thread t @@ fun (Per_thread p) ->
   Mutex.lock t.mutex;
   if t.run then begin
     Mutex.unlock t.mutex;
@@ -401,8 +400,9 @@ let run_fiber ?context:t_opt fiber main =
     t.run <- true;
     Mutex.unlock t.mutex
   end;
-  Mpmcq.push p.ready (Spawn (fiber, main));
-  next pt
+  p.remaining_quota <- t.quota;
+  p.fiber <- Fiber.Maybe.of_fiber fiber;
+  Effect.Deep.match_with main fiber t.handler
 
 let run ?context ?(forbid = false) main =
   let computation = Computation.create ~mode:`LIFO () in
diff --git a/lib/picos_mux.random/picos_mux_random.ml b/lib/picos_mux.random/picos_mux_random.ml
@@ -79,25 +79,25 @@ let[@inline] relaxed_wakeup t ~known_not_empty =
     Condition.signal t.condition
   end
 
-let exec p t = function
-  | Spawn (fiber, main) ->
-      p := Fiber.Maybe.of_fiber fiber;
-      Effect.Deep.match_with main fiber t.handler
-  | Raise (fiber, k, exn, bt) ->
-      p := Fiber.Maybe.of_fiber fiber;
-      Effect.Deep.discontinue_with_backtrace k exn bt
-  | Return (fiber, k) ->
-      p := Fiber.Maybe.of_fiber fiber;
-      Effect.Deep.continue k ()
-  | Current (fiber, k) ->
-      p := Fiber.Maybe.of_fiber fiber;
-      Effect.Deep.continue k fiber
-  | Continue (fiber, k) ->
-      p := Fiber.Maybe.of_fiber fiber;
-      Fiber.continue fiber k ()
-  | Resume (fiber, k) ->
-      p := Fiber.Maybe.of_fiber fiber;
-      Fiber.resume fiber k
+let exec p t ready =
+  begin
+    p :=
+      match ready with
+      | Spawn (fiber, _)
+      | Raise (fiber, _, _, _)
+      | Return (fiber, _)
+      | Current (fiber, _)
+      | Continue (fiber, _)
+      | Resume (fiber, _) ->
+          Fiber.Maybe.of_fiber fiber
+  end;
+  match ready with
+  | Spawn (fiber, main) -> Effect.Deep.match_with main fiber t.handler
+  | Raise (_, k, exn, bt) -> Effect.Deep.discontinue_with_backtrace k exn bt
+  | Return (_, k) -> Effect.Deep.continue k ()
+  | Current (fiber, k) -> Effect.Deep.continue k fiber
+  | Continue (fiber, k) -> Fiber.continue fiber k ()
+  | Resume (fiber, k) -> Fiber.resume fiber k
 
 let rec next p t =
   match Collection.pop_exn t.ready with
@@ -288,8 +288,8 @@ let run_fiber ?context:t_opt fiber main =
   else begin
     t.run <- true;
     Mutex.unlock t.mutex;
-    Collection.push t.ready (Spawn (fiber, main));
-    next (get ()) t;
+    get () := Fiber.Maybe.of_fiber fiber;
+    Effect.Deep.match_with main fiber t.handler;
     Mutex.lock t.mutex;
     await t
   end
