diff --git a/sledge/nonstdlib/NSMap.ml b/sledge/nonstdlib/NSMap.ml index 1dd3f9014..dd24084ec 100644 --- a/sledge/nonstdlib/NSMap.ml +++ b/sledge/nonstdlib/NSMap.ml @@ -148,6 +148,18 @@ struct let existsi m ~f = M.exists (fun key data -> f ~key ~data) m let for_alli m ~f = M.for_all (fun key data -> f ~key ~data) m let fold m s ~f = M.fold (fun key data acc -> f ~key ~data acc) m s + + let fold_until (type res) m s ~f ~finish = + let state = ref s in + let exception Stop of res in + try + iteri m ~f:(fun ~key ~data -> + match f ~key ~data !state with + | `Continue s -> state := s + | `Stop r -> raise_notrace (Stop r) ) ; + finish !state + with Stop r -> r + let keys m = Iter.from_iter (fun f -> M.iter (fun k _ -> f k) m) let values m = Iter.from_iter (fun f -> M.iter (fun _ v -> f v) m) let to_iter m = Iter.from_iter (fun f -> M.iter (fun k v -> f (k, v)) m) diff --git a/sledge/nonstdlib/NSMap_intf.ml b/sledge/nonstdlib/NSMap_intf.ml index 02b594ebe..4261452d5 100644 --- a/sledge/nonstdlib/NSMap_intf.ml +++ b/sledge/nonstdlib/NSMap_intf.ml @@ -122,6 +122,13 @@ module type S = sig val for_alli : 'a t -> f:(key:key -> data:'a -> bool) -> bool val fold : 'a t -> 's -> f:(key:key -> data:'a -> 's -> 's) -> 's + val fold_until : + 'a t + -> 's + -> f:(key:key -> data:'a -> 's -> [`Continue of 's | `Stop of 'b]) + -> finish:('s -> 'b) + -> 'b + (** {1 Convert} *) val keys : 'a t -> key iter