From a4caa0bd65cd009a067a7022b63798a43477dd7e Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Sun, 21 Feb 2021 13:18:34 -0800
Subject: [PATCH] [sledge] Optimize Map operations

Reviewed By: jvillard

Differential Revision: D26338012

fbshipit-source-id: aa52307f1
---
 sledge/nonstdlib/NSMap.ml         | 67 +++----------------------------
 sledge/nonstdlib/NSMap_intf.ml    |  8 ----
 sledge/nonstdlib/multiset.ml      |  3 --
 sledge/nonstdlib/multiset_intf.ml |  9 -----
 sledge/nonstdlib/ocaml/map.ml     | 42 ++++++++++++++++---
 sledge/nonstdlib/ocaml/map.mli    | 30 +++++++++++++-
 6 files changed, 69 insertions(+), 90 deletions(-)

diff --git a/sledge/nonstdlib/NSMap.ml b/sledge/nonstdlib/NSMap.ml
index 812148553..50a2497ec 100644
--- a/sledge/nonstdlib/NSMap.ml
+++ b/sledge/nonstdlib/NSMap.ml
@@ -102,67 +102,12 @@ struct
       m (empty, empty)
 
   let is_empty = M.is_empty
-
-  let root_key m =
-    if M.is_empty m then None
-    else
-      let exception Found in
-      let found = ref None in
-      try
-        M.find_first
-          (fun key ->
-            found := Some key ;
-            raise_notrace Found )
-          m
-        |> ignore ;
-        None
-      with
-      | Found -> !found
-      | Not_found -> None
-
-  let root_binding m =
-    let exception Found in
-    let found = ref None in
-    try
-      M.for_all
-        (fun key data ->
-          found := Some (key, data) ;
-          raise_notrace Found )
-        m
-      |> ignore ;
-      None
-    with
-    | Found -> !found
-    | Not_found -> None
-
-  let is_singleton m =
-    match root_key m with
-    | Some k ->
-        let l, _, r = M.split k m in
-        is_empty l && is_empty r
-    | None -> false
-
+  let is_singleton = M.is_singleton
   let length = M.cardinal
-
-  let only_binding m =
-    match root_key m with
-    | Some k -> (
-      match M.split k m with
-      | l, Some v, r when is_empty l && is_empty r -> Some (k, v)
-      | _ -> None )
-    | None -> None
-
-  let classify m =
-    match root_key m with
-    | None -> Zero2
-    | Some k -> (
-      match M.split k m with
-      | l, Some v, r when is_empty l && is_empty r -> One2 (k, v)
-      | _ -> Many2 )
-
-  let choose_key = root_key
-  let choose = root_binding
-  let choose_exn m = Option.get_exn (choose m)
+  let only_binding = M.only_binding
+  let classify = M.classify
+  let choose = M.choose_opt
+  let choose_exn = M.choose
   let min_binding = M.min_binding_opt
   let mem k m = M.mem k m
   let find_exn k m = M.find k m
@@ -195,8 +140,6 @@ struct
     in
     match !found with Some v -> `Found v | None -> `Added m
 
-  let pop m = choose m |> Option.map ~f:(fun (k, v) -> (k, v, remove k m))
-
   let pop_min_binding m =
     min_binding m |> Option.map ~f:(fun (k, v) -> (k, v, remove k m))
 
diff --git a/sledge/nonstdlib/NSMap_intf.ml b/sledge/nonstdlib/NSMap_intf.ml
index 2e2b4992a..5e2fc68d1 100644
--- a/sledge/nonstdlib/NSMap_intf.ml
+++ b/sledge/nonstdlib/NSMap_intf.ml
@@ -71,10 +71,6 @@ module type S = sig
   val only_binding : 'a t -> (key * 'a) option
   val classify : 'a t -> (key, 'a) zero_one_many2
 
-  val choose_key : 'a t -> key option
-  (** Find an unspecified key. Different keys may be chosen for equivalent
-      maps. [O(1)]. *)
-
   val choose : 'a t -> (key * 'a) option
   (** Find an unspecified binding. Different bindings may be chosen for
       equivalent maps. [O(1)]. *)
@@ -96,10 +92,6 @@ module type S = sig
   (** Find the value bound to the given key if there is one, or otherwise
       add a binding for the given key and value. *)
 
-  val pop : 'a t -> (key * 'a * 'a t) option
-  (** Find and remove an unspecified binding. Different bindings may be
-      chosen for equivalent maps. [O(1)]. *)
-
   val pop_min_binding : 'a t -> (key * 'a * 'a t) option
   (** Find and remove binding with minimum key. [O(log n)]. *)
 
diff --git a/sledge/nonstdlib/multiset.ml b/sledge/nonstdlib/multiset.ml
index 44daada33..315cd44b4 100644
--- a/sledge/nonstdlib/multiset.ml
+++ b/sledge/nonstdlib/multiset.ml
@@ -109,9 +109,6 @@ struct
   let count x m = match M.find x m with Some q -> q | None -> Mul.zero
   let only_elt = M.only_binding
   let classify = M.classify
-  let choose = M.choose
-  let choose_exn = M.choose_exn
-  let pop = M.pop
   let min_elt = M.min_binding
   let pop_min_elt = M.pop_min_binding
   let to_iter = M.to_iter
diff --git a/sledge/nonstdlib/multiset_intf.ml b/sledge/nonstdlib/multiset_intf.ml
index 16689a28b..df1845d17 100644
--- a/sledge/nonstdlib/multiset_intf.ml
+++ b/sledge/nonstdlib/multiset_intf.ml
@@ -94,15 +94,6 @@ module type S = sig
   val only_elt : t -> (elt * mul) option
   (** The only element of a singleton multiset. [O(1)]. *)
 
-  val choose_exn : t -> elt * mul
-  (** Find an unspecified element. [O(1)]. *)
-
-  val choose : t -> (elt * mul) option
-  (** Find an unspecified element. [O(1)]. *)
-
-  val pop : t -> (elt * mul * t) option
-  (** Find and remove an unspecified element. [O(1)]. *)
-
   val min_elt : t -> (elt * mul) option
   (** Minimum element. [O(log n)]. *)
 
diff --git a/sledge/nonstdlib/ocaml/map.ml b/sledge/nonstdlib/ocaml/map.ml
index edc2ef0e6..ceaf06ba1 100644
--- a/sledge/nonstdlib/ocaml/map.ml
+++ b/sledge/nonstdlib/ocaml/map.ml
@@ -13,6 +13,8 @@
 (*                                                                        *)
 (**************************************************************************)
 
+open! NS0
+
 module type OrderedType =
   sig
     type t
@@ -32,6 +34,7 @@ module type S =
     val add: key -> 'a -> 'a t -> 'a t
     val update: key -> ('a option -> 'a option) -> 'a t -> 'a t
     val singleton: key -> 'a -> 'a t
+    val is_singleton: 'a t -> bool
     val remove: key -> 'a t -> 'a t
     val merge:
           (key -> 'a option -> 'b option -> 'c option) -> 'a t -> 'b t -> 'c t
@@ -54,12 +57,16 @@ module type S =
     val partition: (key -> 'a -> bool) -> 'a t -> 'a t * 'a t
     val cardinal: 'a t -> int
     val bindings: 'a t -> (key * 'a) list
+    val only_binding: 'a t -> (key * 'a) option
+    val classify : 'a t -> (key, 'a) zero_one_many2
     val min_binding: 'a t -> (key * 'a)
     val min_binding_opt: 'a t -> (key * 'a) option
     val max_binding: 'a t -> (key * 'a)
     val max_binding_opt: 'a t -> (key * 'a) option
     val choose: 'a t -> (key * 'a)
     val choose_opt: 'a t -> (key * 'a) option
+    val divide : 'a t -> ('a t * key * 'a * 'a t) option
+    val divide_exn : 'a t -> ('a t * key * 'a * 'a t)
     val split: key -> 'a t -> 'a t * 'a option * 'a t
     val find: key -> 'a t -> 'a
     val find_opt: key -> 'a t -> 'a option
@@ -213,8 +220,14 @@ module Make (Ord : Comparer.S) = struct
         t_of_sexp Key.t_of_sexp data_of_sexp Ord.compare_of_sexp s
     end
 
+    let empty = Empty
+
+    let is_empty = function Empty -> true | _ -> false
+
     let singleton x d = Node{l=Empty; v=x; d; r=Empty; h=1}
 
+    let is_singleton = function Node {l=Empty; r=Empty} -> true | _ -> false
+
     let bal l x d r =
       let hl = match l with Empty -> 0 | Node {h} -> h in
       let hr = match r with Empty -> 0 | Node {h} -> h in
@@ -245,10 +258,6 @@ module Make (Ord : Comparer.S) = struct
       end else
         Node{l; v=x; d; r; h=(if hl >= hr then hl + 1 else hr + 1)}
 
-    let empty = Empty
-
-    let is_empty = function Empty -> true | _ -> false
-
     let rec add x data = function
         Empty ->
           Node{l=Empty; v=x; d=data; r=Empty; h=1}
@@ -358,6 +367,15 @@ module Make (Ord : Comparer.S) = struct
           let c = Ord.compare x v in
           c = 0 || mem x (if c < 0 then l else r)
 
+    let classify = function
+      | Empty -> Zero2
+      | Node {l=Empty; v; d; r=Empty} -> One2 (v, d)
+      | _ -> Many2
+
+    let only_binding = function
+        Node {l=Empty; v; d; r=Empty} -> Some (v, d)
+      | _ -> None
+
     let rec min_binding = function
         Empty -> raise Not_found
       | Node {l=Empty; v; d} -> (v, d)
@@ -507,6 +525,14 @@ module Make (Ord : Comparer.S) = struct
       | Some d -> join t1 v d t2
       | None -> concat t1 t2
 
+    let divide_exn = function
+      | Node {l; v; d; r} -> (l, v, d, r)
+      | Empty -> raise Not_found
+
+    let divide = function
+      | Node {l; v; d; r} -> Some (l, v, d, r)
+      | Empty -> None
+
     let rec split x = function
         Empty ->
           (Empty, None, Empty)
@@ -587,9 +613,13 @@ module Make (Ord : Comparer.S) = struct
 
     let bindings = bindings
 
-    let choose = min_binding
+    let choose = function
+        Empty -> raise Not_found
+      | Node {v; d} -> (v, d)
 
-    let choose_opt = min_binding_opt
+    let choose_opt = function
+        Empty -> None
+      | Node {v; d} -> Some (v, d)
 
     let add_seq i m =
       Seq.fold_left (fun m (k,v) -> add k v m) m i
diff --git a/sledge/nonstdlib/ocaml/map.mli b/sledge/nonstdlib/ocaml/map.mli
index c7ecb77e0..66dc94b38 100644
--- a/sledge/nonstdlib/ocaml/map.mli
+++ b/sledge/nonstdlib/ocaml/map.mli
@@ -43,6 +43,8 @@
    values so its type is [string PairsMap.t].
 *)
 
+open! NS0
+
 module type OrderedType =
   sig
     type t
@@ -107,6 +109,9 @@ module type S =
         @since 3.12.0
      *)
 
+    val is_singleton: 'a t -> bool
+    (** Test whether a map contains only a single binding or not. *)
+
     val remove: key -> 'a t -> 'a t
     (** [remove x m] returns a map containing the same bindings as
        [m], except for [x] which is unbound in the returned map.
@@ -226,6 +231,11 @@ module type S =
         @since 3.12.0
      *)
 
+    val only_binding: 'a t -> (key * 'a) option
+    (** Return the binding of a singleton map, or None otherwise. *)
+
+    val classify : 'a t -> (key, 'a) zero_one_many2
+
     val min_binding: 'a t -> (key * 'a)
     (** Return the binding with the smallest key in a given map
        (with respect to the [Ord.compare] ordering), or raise
@@ -255,17 +265,33 @@ module type S =
     val choose: 'a t -> (key * 'a)
     (** Return one binding of the given map, or raise [Not_found] if
        the map is empty. Which binding is chosen is unspecified,
-       but equal bindings will be chosen for equal maps.
+       and different bindings may be chosen for equal maps.
         @since 3.12.0
      *)
 
     val choose_opt: 'a t -> (key * 'a) option
     (** Return one binding of the given map, or [None] if
        the map is empty. Which binding is chosen is unspecified,
-       but equal bindings will be chosen for equal maps.
+       and different bindings may be chosen for equal maps.
         @since 4.05
      *)
 
+    val divide : 'a t -> ('a t * key * 'a * 'a t) option
+    (** [divide m] returns [None] if [m] is empty and otherwise
+        [Some (l, key, data, r)], where
+          [key] is some key bound in [m];
+          [data] is the data associated to [key] in [m];
+          [l] is the map with all the bindings of [m] whose key
+        is strictly less than [key];
+          [r] is the map with all the bindings of [m] whose key
+        is strictly greater than [key].
+        Runs in constant time, and the [l] and [r] maps are close
+        to the same size.
+     *)
+
+    val divide_exn : 'a t -> ('a t * key * 'a * 'a t)
+    (** Same as {!Map.S.divide}, but raises [Not_found] if the map is empty. *)
+
     val split: key -> 'a t -> 'a t * 'a option * 'a t
     (** [split x m] returns a triple [(l, data, r)], where
           [l] is the map with all the bindings of [m] whose key