diff --git a/sledge/lib/import/map.ml b/sledge/lib/import/map.ml index 1c826dc7f..f6d6aeb95 100644 --- a/sledge/lib/import/map.ml +++ b/sledge/lib/import/map.ml @@ -28,6 +28,16 @@ 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 choose m = + with_return + @@ fun {return} -> + binary_search_segmented m `Last_on_left ~segment_of:(fun ~key ~data -> + return (Some (key, data)) ) + |> ignore ; + None + + let pop m = choose m |> Option.map ~f:(fun (k, v) -> (k, v, remove m k)) + let find_and_remove m k = let found = ref None in let m = diff --git a/sledge/lib/import/map_intf.ml b/sledge/lib/import/map_intf.ml index d21cf6de1..3bddeb5fb 100644 --- a/sledge/lib/import/map_intf.ml +++ b/sledge/lib/import/map_intf.ml @@ -23,6 +23,8 @@ module type S = sig val merge_skewed : 'a t -> 'a t -> combine:(key:key -> 'a -> 'a -> 'a) -> 'a t + 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 val pp : key pp -> 'a pp -> 'a t pp diff --git a/sledge/lib/import/qset.ml b/sledge/lib/import/qset.ml index a8fa5ef02..11772e6a5 100644 --- a/sledge/lib/import/qset.ml +++ b/sledge/lib/import/qset.ml @@ -73,6 +73,8 @@ struct let map_counts m ~f = M.mapi ~f:(fun ~key ~data -> f key data) m let length m = M.length m let count m x = match M.find m x with Some q -> q | None -> Q.zero + let choose = M.choose + let pop = M.pop let min_elt_exn = M.min_elt_exn let min_elt = M.min_elt let to_list m = M.to_alist m diff --git a/sledge/lib/import/qset_intf.ml b/sledge/lib/import/qset_intf.ml index a61d25816..8c01787ad 100644 --- a/sledge/lib/import/qset_intf.ml +++ b/sledge/lib/import/qset_intf.ml @@ -49,6 +49,9 @@ module type S = sig val count : t -> elt -> Q.t (** Multiplicity of an element. [O(log n)]. *) + val choose : t -> (elt * Q.t) option + val pop : t -> (elt * Q.t * t) option + val min_elt_exn : t -> elt * Q.t (** Minimum element. *)