From b35e3d30e087de1704b96ae261a6d021c9e3f976 Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Thu, 16 Apr 2020 03:37:49 -0700
Subject: [PATCH] [sledge] Add Map.fold_until

Reviewed By: jvillard

Differential Revision: D20863522

fbshipit-source-id: db4ad2cdb
---
 sledge/lib/import/map.ml      | 8 ++++++++
 sledge/lib/import/map_intf.ml | 7 +++++++
 2 files changed, 15 insertions(+)

diff --git a/sledge/lib/import/map.ml b/sledge/lib/import/map.ml
index f6d6aeb95..876d74ca2 100644
--- a/sledge/lib/import/map.ml
+++ b/sledge/lib/import/map.ml
@@ -28,6 +28,14 @@ end) : S with type key = Key.t = struct
   let merge_skewed x y ~combine =
     of_map (Core.Map.merge_skewed (to_map x) (to_map y) ~combine)
 
+  let fold_until m ~init ~f ~finish =
+    let fold m ~init ~f =
+      let f ~key ~data s = f s (key, data) in
+      fold m ~init ~f
+    in
+    let f s (k, v) = f ~key:k ~data:v s in
+    Container.fold_until ~fold ~init ~f ~finish m
+
   let choose m =
     with_return
     @@ fun {return} ->
diff --git a/sledge/lib/import/map_intf.ml b/sledge/lib/import/map_intf.ml
index 3bddeb5fb..e6e68c06f 100644
--- a/sledge/lib/import/map_intf.ml
+++ b/sledge/lib/import/map_intf.ml
@@ -23,6 +23,13 @@ module type S = sig
   val merge_skewed :
     'a t -> 'a t -> combine:(key:key -> 'a -> 'a -> 'a) -> 'a t
 
+  val fold_until :
+       'v t
+    -> init:'a
+    -> f:(key:key -> data:'v -> 'a -> ('a, 'r) Continue_or_stop.t)
+    -> finish:('a -> 'r)
+    -> 'r
+
   val choose : 'a t -> (key * 'a) option
   val pop : 'a t -> (key * 'a * 'a t) option
   val find_and_remove : 'a t -> key -> ('a * 'a t) option