From 0d3381bef42ae1d0975c2ba2124f04eaf3d01e27 Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Sun, 15 Dec 2024 12:02:12 +0200
Subject: [PATCH] Add `bind_on` and `await_on` to `Picos_lwt`

---
 lib/picos_lwt.unix/picos_lwt_unix.ml  |  2 +-
 lib/picos_lwt.unix/picos_lwt_unix.mli |  3 +++
 lib/picos_lwt/picos_lwt.ml            | 31 +++++++++++++++++++++++++++
 lib/picos_lwt/picos_lwt.mli           |  6 ++++++
 4 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/lib/picos_lwt.unix/picos_lwt_unix.ml b/lib/picos_lwt.unix/picos_lwt_unix.ml
index 19e7faccd..30fc84d05 100644
--- a/lib/picos_lwt.unix/picos_lwt_unix.ml
+++ b/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
diff --git a/lib/picos_lwt.unix/picos_lwt_unix.mli b/lib/picos_lwt.unix/picos_lwt_unix.mli
index 71a9817b0..586e30b66 100644
--- a/lib/picos_lwt.unix/picos_lwt_unix.mli
+++ b/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
diff --git a/lib/picos_lwt/picos_lwt.ml b/lib/picos_lwt/picos_lwt.ml
index 3846f3f92..0154d1ef2 100644
--- a/lib/picos_lwt/picos_lwt.ml
+++ b/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 ->
diff --git a/lib/picos_lwt/picos_lwt.mli b/lib/picos_lwt/picos_lwt.mli
index d6a0d73b2..1e6c626ef 100644
--- a/lib/picos_lwt/picos_lwt.mli
+++ b/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