Repository: ocaml-multicore/picos
Pull request #118 (2024-04-27 12:21:19)

Defunctorize `Picos_lwt`
I realized there is no need to use a functor.  `await` could be just moved outside of the functor and `sleep` is a monomorphic function, which can easily be passed around - a functor is overkill.


Patch:
diff --git a/lib/picos_lwt/picos_lwt.ml b/lib/picos_lwt/picos_lwt.ml
index 0ce0ea8..95d84d7 100644
--- a/lib/picos_lwt/picos_lwt.ml
+++ b/lib/picos_lwt/picos_lwt.ml
@@ -1,98 +1,102 @@
-include Intf
+open Picos
+open Lwt.Infix
 
-module Make (Sleep : Sleep) : S = struct
-  open Picos
-  open Lwt.Infix
+let await thunk =
+  let computation = Computation.create () in
+  let promise =
+    Lwt.try_bind thunk
+      (fun value ->
+        Computation.return computation value;
+        Lwt.return_unit)
+      (fun exn ->
+        Computation.cancel computation (Exn_bt.get_callstack 0 exn);
+        Lwt.return_unit)
+  in
+  Lwt.async (fun () -> promise);
+  let trigger = Trigger.create () in
+  if Computation.try_attach computation trigger then begin
+    match Trigger.await trigger with
+    | None -> Computation.await computation
+    | Some exn_bt ->
+        Lwt.cancel promise;
+        Exn_bt.raise exn_bt
+  end
+  else Computation.await computation
 
-  let[@alert "-handler"] rec run :
-      type a r.
-      Fiber.t ->
-      (a, r) Effect.Shallow.continuation ->
-      (a, Exn_bt.t) Result.t ->
-      r Lwt.t =
-   fun fiber k v ->
-    let effc (type a) :
-        a Effect.t -> ((a, _) Effect.Shallow.continuation -> _) option =
-      function
-      | Fiber.Current -> Some (fun k -> run fiber k (Ok fiber))
-      | Fiber.Spawn r ->
-          Some
-            (fun k ->
-              match Fiber.canceled fiber with
-              | None ->
-                  let packed = Computation.Packed r.computation in
-                  List.iter
-                    (fun main ->
-                      let fiber = Fiber.create_packed ~forbid:r.forbid packed in
-                      Lwt.async @@ fun () ->
-                      run fiber (Effect.Shallow.fiber main) (Ok ()))
-                    r.mains;
-                  run fiber k (Ok ())
-              | Some exn_bt -> run fiber k (Error exn_bt))
-      | Fiber.Yield ->
-          Some
-            (fun k ->
-              match Fiber.canceled fiber with
-              | None -> Lwt.pause () >>= fun () -> run fiber k (Ok ())
-              | Some exn_bt -> run fiber k (Error exn_bt))
-      | Computation.Cancel_after r ->
-          Some
-            (fun k ->
-              match Fiber.canceled fiber with
-              | None ->
-                  let sleep =
-                    Sleep.sleep r.seconds >>= fun () ->
-                    Computation.cancel r.computation r.exn_bt;
-                    Lwt.return_unit
-                  in
-                  let canceler =
-                    Trigger.from_action sleep () @@ fun _ sleep _ ->
-                    Lwt.cancel sleep
-                  in
-                  if Computation.try_attach r.computation canceler then
-                    Lwt.async @@ fun () -> sleep
-                  else Trigger.signal canceler;
-                  run fiber k (Ok ())
-              | Some exn_bt -> run fiber k (Error exn_bt))
-      | Trigger.Await trigger ->
-          Some
-            (fun k ->
-              let promise, resolver = Lwt.wait () in
-              let resume _trigger resolver _ = Lwt.wakeup resolver () in
-              if Fiber.try_suspend fiber trigger resolver () resume then
-                promise >>= fun () -> run fiber k (Ok (Fiber.canceled fiber))
-              else run fiber k (Ok (Fiber.canceled fiber)))
-      | _ -> None
-    in
-    let handler = Effect.Shallow.{ retc = Lwt.return; exnc = Lwt.fail; effc } in
-    match v with
-    | Ok v -> Effect.Shallow.continue_with k v handler
-    | Error exn_bt -> Exn_bt.discontinue_with k exn_bt handler
+let[@alert "-handler"] rec run :
+    type a r.
+    Fiber.t ->
+    (float -> unit Lwt.t) ->
+    (a, r) Effect.Shallow.continuation ->
+    (a, Exn_bt.t) Result.t ->
+    r Lwt.t =
+ fun fiber sleep k v ->
+  let effc (type a) :
+      a Effect.t -> ((a, _) Effect.Shallow.continuation -> _) option = function
+    | Fiber.Current -> Some (fun k -> run fiber sleep k (Ok fiber))
+    | Fiber.Spawn r ->
+        Some
+          (fun k ->
+            match Fiber.canceled fiber with
+            | None ->
+                let packed = Computation.Packed r.computation in
+                List.iter
+                  (fun main ->
+                    let fiber = Fiber.create_packed ~forbid:r.forbid packed in
+                    Lwt.async @@ fun () ->
+                    run fiber sleep (Effect.Shallow.fiber main) (Ok ()))
+                  r.mains;
+                run fiber sleep k (Ok ())
+            | Some exn_bt -> run fiber sleep k (Error exn_bt))
+    | Fiber.Yield ->
+        Some
+          (fun k ->
+            match Fiber.canceled fiber with
+            | None -> Lwt.pause () >>= fun () -> run fiber sleep k (Ok ())
+            | Some exn_bt -> run fiber sleep k (Error exn_bt))
+    | Computation.Cancel_after r ->
+        Some
+          (fun k ->
+            match Fiber.canceled fiber with
+            | None ->
+                let timeout =
+                  sleep r.seconds >>= fun () ->
+                  Computation.cancel r.computation r.exn_bt;
+                  Lwt.return_unit
+                in
+                let canceler =
+                  Trigger.from_action timeout () @@ fun _ sleep _ ->
+                  Lwt.cancel sleep
+                in
+                if Computation.try_attach r.computation canceler then
+                  Lwt.async @@ fun () -> timeout
+                else Trigger.signal canceler;
+                run fiber sleep k (Ok ())
+            | Some exn_bt -> run fiber sleep k (Error exn_bt))
+    | Trigger.Await trigger ->
+        Some
+          (fun k ->
+            let promise, resolver = Lwt.wait () in
+            let resume _trigger resolver _ = Lwt.wakeup resolver () in
+            if Fiber.try_suspend fiber trigger resolver () resume then
+              promise >>= fun () ->
+              run fiber sleep k (Ok (Fiber.canceled fiber))
+            else run fiber sleep k (Ok (Fiber.canceled fiber)))
+    | _ -> None
+  in
+  let handler = Effect.Shallow.{ retc = Lwt.return; exnc = Lwt.fail; effc } in
+  match v with
+  | Ok v -> Effect.Shallow.continue_with k v handler
+  | Error exn_bt -> Exn_bt.discontinue_with k exn_bt handler
 
-  let run ?(forbid = false) main =
-    let computation = Computation.create () in
-    let fiber = Fiber.create ~forbid computation in
-    run fiber (Effect.Shallow.fiber main) (Ok ())
+let run ?(forbid = false) ~sleep main =
+  let computation = Computation.create () in
+  let fiber = Fiber.create ~forbid computation in
+  run fiber sleep (Effect.Shallow.fiber main) (Ok ())
 
-  let await thunk =
-    let computation = Computation.create () in
-    let promise =
-      Lwt.try_bind thunk
-        (fun value ->
-          Computation.return computation value;
-          Lwt.return_unit)
-        (fun exn ->
-          Computation.cancel computation (Exn_bt.get_callstack 0 exn);
-          Lwt.return_unit)
-    in
-    Lwt.async (fun () -> promise);
-    let trigger = Trigger.create () in
-    if Computation.try_attach computation trigger then begin
-      match Trigger.await trigger with
-      | None -> Computation.await computation
-      | Some exn_bt ->
-          Lwt.cancel promise;
-          Exn_bt.raise exn_bt
-    end
-    else Computation.await computation
+include Intf
+
+module Make (Sleep : Sleep) : S = struct
+  let run ?forbid main = run ?forbid ~sleep:Sleep.sleep main
+  let await = await
 end
diff --git a/lib/picos_lwt/picos_lwt.mli b/lib/picos_lwt/picos_lwt.mli
index 4062693..36f8055 100644
--- a/lib/picos_lwt/picos_lwt.mli
+++ b/lib/picos_lwt/picos_lwt.mli
@@ -1,10 +1,36 @@
-(** A functor for building a {!Picos} compatible direct style interface to
-    {!Lwt} with given implementation of {{!Sleep} sleep}.
+(** A {!Picos} compatible direct style interface to {!Lwt} with given
+    implementation of [sleep].
 
     This basically gives you an alternative direct style interface to
     programming with {!Lwt}.  All the scheduling decisions will be made by
     {!Lwt}. *)
 
+val await : (unit -> 'a Lwt.t) -> 'a
+(** [await thunk] awaits for the promise returned by [thunk ()] to resolve and
+    returns the result.  This should only be called from inside a fiber started
+    through {!run}. *)
+
+val run :
+  ?forbid:bool -> sleep:(float -> unit Lwt.t) -> (unit -> 'a) -> 'a Lwt.t
+(** [run ~sleep main] runs the [main] program implemented in {!Picos} as a
+    promise with {!Lwt} as the scheduler and given operation to [sleep].  In
+    other words, the [main] program will be run as a {!Lwt} promise or fiber.
+
+    Calling [sleep seconds] should return a cancelable promise that resolves
+    after given number of [seconds] (unless canceled).
+
+    ℹ️ Inside [main] you can use anything implemented in Picos for concurrent
+    programming.  In particular, you only need to call [run] with an
+    implementation of [sleep] at the entry point of your application.
+
+    The optional [forbid] argument defaults to [false] and determines whether
+    propagation of cancelation is initially allowed. *)
+
+(** {1 Functorized interface}
+
+    A functor for building a {!Picos} compatible direct style interface to
+    {!Lwt} with given implementation of {{!Sleep} sleep}. *)
+
 include module type of Intf
 
 (** [Make (Sleep)] creates a {!Picos} compatible interface to {!Lwt} with given
@@ -17,3 +43,5 @@ include module type of Intf
     instantiates this functor using {!Lwt_unix.sleep} as the implemention of
     {{!Sleep.sleep} sleep}. *)
 module Make : functor (_ : Sleep) -> S
+[@@deprecated
+  "Just use Picos_lwt.run instead. This functorized interface will be removed."]


Test patch:
diff --git a/test/test_lwt_unix.ml b/test/test_lwt_unix.ml
index 358d118..05c5ed2 100644
--- a/test/test_lwt_unix.ml
+++ b/test/test_lwt_unix.ml
@@ -1,14 +1,14 @@
 open Picos
-module Picos_lwt_unix = Picos_lwt.Make (Lwt_unix)
 
 let basics () =
-  Lwt_main.run @@ Picos_lwt_unix.run
+  Lwt_main.run
+  @@ Picos_lwt.run ~sleep:Lwt_unix.sleep
   @@ fun () ->
   let computation = Computation.create () in
   let child =
     Computation.capture computation @@ fun () ->
     while true do
-      Picos_lwt_unix.await (fun () -> Lwt_unix.sleep 0.01)
+      Picos_lwt.await (fun () -> Lwt_unix.sleep 0.01)
     done
   in
   Fiber.spawn ~forbid:false computation [ child ];