Add bind_on and await_on to Picos_lwt

Author: polytypic

lib/picos_lwt.unix/picos_lwt_unix.ml

@@ -41,7 +41,7 @@ let rec notify () =
       end
       else notify ()
 
-module System = struct
+module System : Picos_lwt.System = struct
   let sleep = Lwt_unix.sleep
 
   type trigger = unit Lwt.t * unit Lwt.u

lib/picos_lwt.unix/picos_lwt_unix.mli

@@ -3,6 +3,9 @@
 
 open Picos
 
+val system : (module Picos_lwt.System)
+(** The system module for Unix. *)
+
 val run_fiber : Fiber.t -> (Fiber.t -> unit) -> unit Lwt.t
 (** [run_fiber fiber main] runs the [main] program as the specified [fiber] as a
     promise with {!Lwt} as the scheduler using a {!Lwt_unix} based

lib/picos_lwt/picos_lwt.ml

@@ -35,6 +35,37 @@ let await promise =
   | Return value -> value
   | Fail exn -> raise exn
 
+let bind_on (module System : System) thunk =
+  let trigger = System.trigger () in
+  let promise = Lwt.bind (System.await trigger) thunk in
+  System.signal trigger;
+  promise
+
+let await_on (module System : System) promise =
+  let computation = Computation.create ~mode:`LIFO () in
+  let trigger = System.trigger () in
+  let promise =
+    Lwt.bind (System.await trigger) @@ fun () ->
+    Lwt.try_bind
+      (fun () -> promise)
+      (fun value ->
+        Computation.return computation value;
+        Lwt.return_unit)
+      (fun exn ->
+        Computation.cancel computation exn empty_bt;
+        Lwt.return_unit)
+  in
+  System.signal trigger;
+  let trigger = Trigger.create () in
+  if Computation.try_attach computation trigger then begin
+    match Trigger.await trigger with
+    | None -> Computation.peek_exn computation
+    | Some (exn, bt) ->
+        Lwt.cancel promise;
+        Printexc.raise_with_backtrace exn bt
+  end
+  else Computation.peek_exn computation
+
 let[@alert "-handler"] rec go : type a r.
     Fiber.t ->
     (module System) ->

lib/picos_lwt/picos_lwt.mli

@@ -17,6 +17,12 @@ val await : 'a Lwt.t -> 'a
 
 include module type of Intf
 
+val bind_on : (module System) -> (unit -> 'a Lwt.t) -> 'a Lwt.t
+(** *)
+
+val await_on : (module System) -> 'a Lwt.t -> 'a
+(** *)
+
 val run_fiber : (module System) -> Fiber.t -> (Fiber.t -> unit) -> unit Lwt.t
 (** [run_fiber (module System) fiber main] runs the [main] program as the
     specified [fiber] as a promise with {!Lwt} as the scheduler using the given