[sledge] Add Map.fold_until

Reviewed By: jvillard

Differential Revision: D20863522

fbshipit-source-id: db4ad2cdb

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} ->

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