[stdune] Add Map.interruptible_fold

The implementation is quite hackish and should be upstreamed;
suggestions welcome.

Signed-off-by: Emilio Jesus Gallego Arias <[email protected]>

src/stdune/map.ml

@@ -58,6 +58,26 @@ module Make (Key : Key) : S with type key = Key.t = struct
 
   let fold t ~init ~f = foldi t ~init ~f:(fun _ x acc -> f x acc)
 
+  type ('a, 'b) interruptible_fold_step =
+    | Continue of 'a
+    | Break of 'b
+
+  let interruptible_fold (t : 'a t) ~(init : 'b) ~(f: 'a -> 'b -> ('b, 'c) interruptible_fold_step) : ('b, 'c) interruptible_fold_step =
+    let exception Break_fold in
+    let break_data = ref None in
+    let f data acc =
+      match f data acc with
+      | Continue acc -> acc
+      | Break res ->
+        break_data := Some res;
+        raise Break_fold
+    in
+    try
+      Continue (fold t ~init ~f)
+    with
+    | Break_fold ->
+      Break (Option.value_exn !break_data)
+
   let for_alli t ~f = for_all t ~f
 
   let for_all t ~f = for_alli t ~f:(fun _ x -> f x)

src/stdune/map_intf.ml

@@ -49,6 +49,17 @@ module type S = sig
 
   val foldi : 'a t -> init:'b -> f:(key -> 'a -> 'b -> 'b) -> 'b
 
+  (* At some point this should be moved to [List] *)
+  type ('a, 'b) interruptible_fold_step =
+    | Continue of 'a
+    | Break of 'b
+
+  val interruptible_fold
+    : 'a t
+    -> init:'b
+    -> f:('a -> 'b -> ('b, 'c) interruptible_fold_step)
+    -> ('b, 'c) interruptible_fold_step
+
   val for_all : 'a t -> f:('a -> bool) -> bool
 
   val for_alli : 'a t -> f:(key -> 'a -> bool) -> bool