From 867444d975cc28fd8acac2b655cbc77d5adc544c Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 4 Mar 2024 22:05:47 -0500
Subject: [PATCH] perf: add `Fut.raise_if_failed`, use it in
 Fiber.check_if_cancelled

---
 src/core/fut.ml   |  5 +++++
 src/core/fut.mli  |  4 ++++
 src/fib/fiber.ml  | 17 ++++++++++-------
 src/fib/fiber.mli |  3 ++-
 4 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/src/core/fut.ml b/src/core/fut.ml
index 3993b171..91a95a7f 100644
--- a/src/core/fut.ml
+++ b/src/core/fut.ml
@@ -29,6 +29,11 @@ let[@inline] peek self : _ option =
   | Done x -> Some x
   | Waiting _ -> None
 
+let[@inline] raise_if_failed self : unit =
+  match A.get self.st with
+  | Done (Error ebt) -> Exn_bt.raise ebt
+  | _ -> ()
+
 let[@inline] is_done self : bool =
   match A.get self.st with
   | Done _ -> true
diff --git a/src/core/fut.mli b/src/core/fut.mli
index e3fea3fc..3b4e4ef4 100644
--- a/src/core/fut.mli
+++ b/src/core/fut.mli
@@ -91,6 +91,10 @@ val is_failed : _ t -> bool
 (** Checks if the future is resolved with [Error _] as a result.
     @since NEXT_RELEASE *)
 
+val raise_if_failed : _ t -> unit
+(** [raise_if_failed fut] raises [e] if [fut] failed with [e].
+    @since NEXT_RELEASE *)
+
 (** {2 Combinators} *)
 
 val spawn : on:Runner.t -> (unit -> 'a) -> 'a t
diff --git a/src/fib/fiber.ml b/src/fib/fiber.ml
index fdfd3830..4f776c80 100644
--- a/src/fib/fiber.ml
+++ b/src/fib/fiber.ml
@@ -291,15 +291,18 @@ let with_on_self_cancel cb (k : unit -> 'a) : 'a =
 
 module Suspend_ = Moonpool.Private.Suspend_
 
+let[@inline] check_if_cancelled_ (self : _ t) = Fut.raise_if_failed self.res
+
 let check_if_cancelled () =
   match Task_local_storage.get k_current_fiber with
   | None ->
     failwith "Fiber.check_if_cancelled: must be run from inside a fiber."
-  | Some (Any self) ->
-    (match peek self with
-    | Some (Error ebt) -> Exn_bt.raise ebt
-    | _ -> ())
+  | Some (Any self) -> check_if_cancelled_ self
 
-let[@inline] yield () : unit =
-  check_if_cancelled ();
-  Suspend_.yield ()
+let yield () : unit =
+  match Task_local_storage.get k_current_fiber with
+  | None -> failwith "Fiber.yield: must be run from inside a fiber."
+  | Some (Any self) ->
+    check_if_cancelled_ self;
+    Suspend_.yield ();
+    check_if_cancelled_ self
diff --git a/src/fib/fiber.mli b/src/fib/fiber.mli
index befc7c35..25787cc1 100644
--- a/src/fib/fiber.mli
+++ b/src/fib/fiber.mli
@@ -81,7 +81,8 @@ val wait_block : 'a t -> 'a Fut.or_error
 val check_if_cancelled : unit -> unit
 (** Check if the current fiber is cancelled, in which case this raises.
     Must be run from inside a fiber.
-    @raise Failure if not. *)
+    @raise e if the current fiber is cancelled with exception [e]
+    @raise Failure if not run from a fiber. *)
 
 val yield : unit -> unit
 (** Yield control to the scheduler from the current fiber.