From 4613aafb30a6e49af817d3dfe5946ecbd53228ef Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 7 Feb 2025 16:57:36 -0500
Subject: [PATCH] feat: add `CCFun.with_return`

---
 src/core/CCFun.ml   |  7 +++++++
 src/core/CCFun.mli  | 16 ++++++++++++++++
 tests/core/t_fun.ml | 12 ++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/src/core/CCFun.ml b/src/core/CCFun.ml
index abb2a6a2..96484352 100644
--- a/src/core/CCFun.ml
+++ b/src/core/CCFun.ml
@@ -61,6 +61,13 @@ let rec iterate n f x =
   else
     iterate (n - 1) f (f x)
 
+let[@inline] with_return (type ret) f : ret =
+  let exception E of ret in
+  let return x = raise_notrace (E x) in
+  match f return with
+  | res -> res
+  | exception E res -> res
+
 module Infix = struct
   (* default implem for some operators *)
   let ( %> ) = compose
diff --git a/src/core/CCFun.mli b/src/core/CCFun.mli
index 88261101..f79a0432 100644
--- a/src/core/CCFun.mli
+++ b/src/core/CCFun.mli
@@ -78,6 +78,22 @@ val iterate : int -> ('a -> 'a) -> 'a -> 'a
     [x], [iterate 1 f x] is [f x], [iterate 2 f x] is [f (f x)], etc.
     @since 2.1 *)
 
+val with_return : (('ret -> 'a) -> 'ret) -> 'ret
+(** [with_return f] is [f return], where [return] is a function
+  that can be invoked to exit the scope early.
+
+  For example:
+{[
+let find_array arr x =
+  let@ return = with_return in
+  for i = 0 to Array.length arr-1 do
+    if arr.(i) = x then return i;
+  done;
+  -1
+]}
+
+  @since NEXT_RELEASE *)
+
 (** {2 Infix}
 
     Infix operators. *)
diff --git a/tests/core/t_fun.ml b/tests/core/t_fun.ml
index 48d64716..587f5c63 100644
--- a/tests/core/t_fun.ml
+++ b/tests/core/t_fun.ml
@@ -21,3 +21,15 @@ true
 t @@ fun () -> CCFun.((succ %> string_of_int) 2 = "3");;
 t @@ fun () -> CCFun.((( * ) 3 % succ) 5 = 18);;
 t @@ fun () -> CCFun.(succ @@ ( * ) 2 @@ pred @@ 3 = 5)
+
+let find_array arr x =
+  let@ return = with_return in
+  for i = 0 to Array.length arr - 1 do
+    if arr.(i) = x then return i
+  done;
+  -1
+;;
+
+eq 1 @@ find_array [| "a"; "b"; "c" |] "b";;
+eq 2 @@ find_array [| "a"; "b"; "c" |] "c";;
+eq (-1) @@ find_array [| "a"; "b"; "c" |] "hello"