perf(fiber): defunctionalize scheduler computations

src/fiber/src/fiber.ml

@@ -25,8 +25,8 @@ type fill = Scheduler.fill = Fill : 'a ivar * 'a -> fill
 module Expert = struct
   type nonrec 'a k = 'a k
 
-  let suspend f k = suspend f k
-  let resume a x k = resume a x k
+  let suspend = suspend
+  let resume = resume
 end
 
 module Temp = Stdune.Temp.Monad (struct

src/fiber/src/lazy.ml

@@ -60,8 +60,15 @@ let is_value t =
   | Running _ | Init _ -> false
 ;;
 
+let stop = End
+
+let run_force t v f k =
+  match eval (execute t v f) stop with
+  | End_of_fiber () -> continue k ()
+  | eff -> Fork (eff, Continue_work k)
+;;
+
 let force_all_unit =
-  let stop () = end_of_fiber in
   (* Fork all computations that haven't been forced yet. Note that this should be
      substantially more efficient that [parallel_map ~f:force] since we ignore
      computations which have already been forced. *)
@@ -71,7 +78,7 @@ let force_all_unit =
       | Done _ | Running _ -> return ()
       | Init f ->
         let v = prep t in
-        fun k -> fork (fun () -> (execute t v f) stop) k)
+        primitive3 run_force t v f)
   in
   (* Wait for all computations, collecting all exceptions.  *)
   (* CR-someday rgrinberg: use [Appendable.t] for [acc] rather than [Appendable.t option]. *)

src/fiber/src/mutex.ml

@@ -7,23 +7,27 @@ type t =
   ; mutable waiters : unit k Queue.t
   }
 
-let lock t k =
+let run_lock t k =
   if t.locked
-  then suspend (fun k -> Queue.push t.waiters k) k
+  then Suspend ((fun k -> Queue.push t.waiters k), k)
   else (
     t.locked <- true;
-    k ())
+    continue k ())
 ;;
 
-let unlock t k =
+let lock t = primitive run_lock t
+
+let run_unlock t k =
   assert t.locked;
   match Queue.pop t.waiters with
   | None ->
     t.locked <- false;
-    k ()
-  | Some next -> resume next () k
+    continue k ()
+  | Some next -> Resume (next, (), k)
 ;;
 
+let unlock t = primitive run_unlock t
+
 let with_lock t ~f =
   let* () = lock t in
   finalize f ~finally:(fun () -> unlock t)

src/fiber/src/mvar.ml

@@ -21,26 +21,30 @@ let create_full x =
   { value = Some x; writers = Queue.create (); readers = Queue.create () }
 ;;
 
-let read t k =
+let run_read t k =
   match t.value with
-  | None -> suspend (fun k -> Queue.push t.readers k) k
+  | None -> Suspend ((fun k -> Queue.push t.readers k), k)
   | Some v ->
     (match Queue.pop t.writers with
      | None ->
        t.value <- None;
-       k v
+       continue k v
      | Some (v', w) ->
        t.value <- Some v';
-       resume w () (fun () -> k v))
+       Resume (w, (), Map ((fun () -> v), k)))
 ;;
 
-let write t x k =
+let read t = primitive run_read t
+
+let run_write t x k =
   match t.value with
-  | Some _ -> suspend (fun k -> Queue.push t.writers (x, k)) k
+  | Some _ -> Suspend ((fun k -> Queue.push t.writers (x, k)), k)
   | None ->
     (match Queue.pop t.readers with
      | None ->
        t.value <- Some x;
-       k ()
-     | Some r -> resume r x (fun () -> k ()))
+       continue k ()
+     | Some r -> Resume (r, x, Map ((fun () -> ()), k)))
 ;;
+
+let write t x = primitive2 run_write t x

src/fiber/src/pool.ml

@@ -25,39 +25,44 @@ type nonrec t =
   ; mutable status : status
   }
 
-let running t k =
+let run_running t k =
   match t.status with
-  | Open -> k true
-  | Closed -> k false
+  | Open -> continue k true
+  | Closed -> continue k false
 ;;
 
+let running t = primitive run_running t
 let create () = { tasks = Queue.create (); runner = Awaiting_run; status = Open }
 
-let task t ~f k =
+let run_task t f k =
   match t.status with
   | Closed -> Code_error.raise "pool is closed. new tasks may not be submitted" []
   | Open ->
     Queue.push t.tasks f;
     (match t.runner with
-     | Running | Awaiting_run -> k ()
+     | Running | Awaiting_run -> continue k ()
      | Awaiting_resume r ->
        t.runner <- Running;
-       resume r () k)
+       Resume (r, (), k))
 ;;
 
-let close t k =
+let task t ~f = primitive2 run_task t f
+
+let run_close t k =
   match t.status with
-  | Closed -> k ()
+  | Closed -> continue k ()
   | Open ->
     t.status <- Closed;
     (match t.runner with
-     | Running | Awaiting_run -> k ()
+     | Running | Awaiting_run -> continue k ()
      | Awaiting_resume r ->
        t.runner <- Running;
-       resume r () k)
+       Resume (r, (), k))
 ;;
 
-let run t k =
+let close t = primitive run_close t
+
+let run_pool t k =
   match t.runner with
   | Awaiting_resume _ | Running ->
     Code_error.raise "Fiber.Pool.run: concurent calls to run aren't allowed" []
@@ -68,14 +73,16 @@ let run t k =
     let n = ref 1 in
     let done_fiber () =
       decr n;
-      if !n = 0 then k () else end_of_fiber
+      if !n = 0 then continue k () else end_of_fiber
     in
     let rec read t =
       match Queue.pop t.tasks with
       | None -> finish_or_suspend t
       | Some v ->
         incr n;
-        fork (fun () -> v () done_fiber) read_delayed
+        (match apply_t v () (Function done_fiber) with
+         | End_of_fiber () -> read_delayed ()
+         | eff -> Fork (eff, Function_work read_delayed))
     and read_delayed () = read t
     and suspend_k k =
       (* we are suspending because we have no tasks *)
@@ -84,11 +91,13 @@ let run t k =
     and finish_or_suspend t =
       match t.status with
       | Closed -> done_fiber ()
-      | Open -> suspend suspend_k read_delayed
+      | Open -> Suspend (suspend_k, Function read_delayed)
     in
     read t
 ;;
 
+let run t = primitive run_pool t
+
 let with_ f =
   of_thunk (fun () ->
     let pool = create () in