Repository: ocaml-multicore/picos
Pull request #157 (2024-06-05 11:29:10)

Drop thunking from `Picos_lwt.await`


Patch:
diff --git a/lib/picos_lwt/picos_lwt.ml b/lib/picos_lwt/picos_lwt.ml
index df8a5ee..ffde6fa 100644
--- a/lib/picos_lwt/picos_lwt.ml
+++ b/lib/picos_lwt/picos_lwt.ml
@@ -5,28 +5,33 @@ include Intf
 let[@inline never] not_main_thread () =
   invalid_arg "not called from the main thread"
 
-let await thunk =
-  if not (Picos_thread.is_main_thread ()) then not_main_thread ();
-  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 await promise =
+  match Lwt.state promise with
+  | Sleep ->
+      if not (Picos_thread.is_main_thread ()) then not_main_thread ();
+      let computation = Computation.create () in
+      let promise =
+        Lwt.try_bind
+          (fun () -> promise)
+          (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
+  | Return value -> value
+  | Fail exn -> raise exn
 
 let[@alert "-handler"] rec go :
     type a r.
diff --git a/lib/picos_lwt/picos_lwt.mli b/lib/picos_lwt/picos_lwt.mli
index 5070c61..5510a81 100644
--- a/lib/picos_lwt/picos_lwt.mli
+++ b/lib/picos_lwt/picos_lwt.mli
@@ -7,9 +7,8 @@
     ℹ️ This is a {{!System} system} independent interface to {!Lwt}.  See
     {!Picos_lwt_unix} for a {!Unix} specific interface. *)
 
-val await : (unit -> 'a Lwt.t) -> 'a
-(** [await thunk] awaits for the promise returned by [thunk ()] to resolve and
-    returns the result.
+val await : 'a Lwt.t -> 'a
+(** [await promise] awaits for the [promise] to resolve and returns the result.
 
     ⚠️ This may only be called on the main thread on which {!Lwt} runs from
     inside a fiber started through {!run}. *)


Test patch:
diff --git a/test/test_lwt_unix.ml b/test/test_lwt_unix.ml
index 9999317..c24eec6 100644
--- a/test/test_lwt_unix.ml
+++ b/test/test_lwt_unix.ml
@@ -7,7 +7,7 @@ let basics () =
   let child =
     Computation.capture computation @@ fun () ->
     while true do
-      Picos_lwt.await (fun () -> Lwt_unix.sleep 0.01)
+      Picos_lwt.await (Lwt_unix.sleep 0.01)
     done
   in
   Fiber.spawn ~forbid:false computation [ child ];