From 3a729f83db5d6df7beeac9374e3a729165d5f942 Mon Sep 17 00:00:00 2001
From: Emilio Jesus Gallego Arias <e+git@x80.org>
Date: Tue, 4 Feb 2020 13:09:15 +0100
Subject: [PATCH] [stdune] Add Map.interruptible_fold

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

Signed-off-by: Emilio Jesus Gallego Arias <e+git@x80.org>
---
 src/stdune/map.ml      | 20 ++++++++++++++++++++
 src/stdune/map_intf.ml | 11 +++++++++++
 2 files changed, 31 insertions(+)

diff --git a/src/stdune/map.ml b/src/stdune/map.ml
index ee6aae220f7d..d7f73f29aad7 100644
--- a/src/stdune/map.ml
+++ b/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)
diff --git a/src/stdune/map_intf.ml b/src/stdune/map_intf.ml
index edc4d4c9ab83..a9fa3e265608 100644
--- a/src/stdune/map_intf.ml
+++ b/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