[sledge] Add Map and Qset operations to access top of underlying tree

Reviewed By: jvillard

Differential Revision: D20663960

fbshipit-source-id: 40e2701c1

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 =

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
 

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

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. *)