[sledge] Optimize Set operations

Reviewed By: jvillard

Differential Revision: D26338013

fbshipit-source-id: 3a21efa9b

sledge/nonstdlib/NSSet.ml

@@ -71,48 +71,12 @@ struct
   let subset s ~of_:t = S.subset s t
   let disjoint = S.disjoint
   let max_elt = S.max_elt_opt
-
-  let root_elt s =
-    let exception Found in
-    let found = ref None in
-    try
-      S.for_all
-        (fun elt ->
-          found := Some elt ;
-          raise_notrace Found )
-        s
-      |> ignore ;
-      None
-    with Found -> !found
-
-  let choose = root_elt
-  let choose_exn m = Option.get_exn (choose m)
-
-  let only_elt s =
-    match root_elt s with
-    | Some elt -> (
-      match S.split elt s with
-      | l, _, r when is_empty l && is_empty r -> Some elt
-      | _ -> None )
-    | None -> None
-
-  let classify s =
-    match root_elt s with
-    | None -> Zero
-    | Some elt -> (
-      match S.split elt s with
-      | l, true, r when is_empty l && is_empty r -> One elt
-      | _ -> Many )
-
-  let pop s =
-    match choose s with
-    | Some elt -> Some (elt, S.remove elt s)
-    | None -> None
-
-  let pop_exn s =
-    let elt = choose_exn s in
-    (elt, S.remove elt s)
-
+  let choose = S.choose_opt
+  let choose_exn = S.choose
+  let pop = S.pop_opt
+  let pop_exn = S.pop
+  let only_elt = S.only_elt
+  let classify = S.classify
   let map s ~f = S.map f s
   let filter s ~f = S.filter f s
   let partition s ~f = S.partition f s

sledge/nonstdlib/NSSet_intf.ml

@@ -74,6 +74,14 @@ module type S = sig
   val only_elt : t -> elt option
   val classify : t -> elt zero_one_many
 
+  val choose : t -> elt option
+  (** Find an unspecified element. Different elements may be chosen for
+      equivalent sets. [O(1)]. *)
+
+  val choose_exn : t -> elt
+  (** Find an unspecified element. Different elements may be chosen for
+      equivalent sets. [O(1)]. *)
+
   val pop : t -> (elt * t) option
   (** Find and remove an unspecified element. [O(1)]. *)
 

sledge/nonstdlib/ocaml/set.ml

@@ -57,12 +57,16 @@ module type S =
     val partition: (elt -> bool) -> t -> t * t
     val cardinal: t -> int
     val elements: t -> elt list
+    val only_elt: t -> elt option
+    val classify : t -> elt NS0.zero_one_many
     val min_elt: t -> elt
     val min_elt_opt: t -> elt option
     val max_elt: t -> elt
     val max_elt_opt: t -> elt option
     val choose: t -> elt
     val choose_opt: t -> elt option
+    val pop : t -> elt * t
+    val pop_opt : t -> (elt * t) option
     val split: elt -> t -> t * bool * t
     val find: elt -> t -> elt
     val find_opt: elt -> t -> elt option
@@ -300,6 +304,16 @@ module Make(Ord: Comparer.S) =
           if rh > lh + 2 then bal (join l v rl) rv rr else
           create l v r
 
+    let classify x : _ NS0.zero_one_many =
+      match x with
+      | Empty -> Zero
+      | Node {l=Empty; v; r=Empty} -> One v
+      | _ -> Many
+
+    let only_elt = function
+        Node {l=Empty; v; r=Empty} -> Some v
+      | _ -> None
+
     (* Smallest and greatest element of a set *)
 
     let rec min_elt = function
@@ -523,9 +537,21 @@ module Make(Ord: Comparer.S) =
 
     let elements = elements
 
-    let choose = min_elt
+    let choose = function
+        Empty -> raise Not_found
+      | Node{v} -> v
+
+    let choose_opt = function
+        Empty -> None
+      | Node{v} -> Some v
+
+    let pop = function
+        Empty -> raise Not_found
+      | Node{l; v; r} -> (v, merge l r)
 
-    let choose_opt = min_elt_opt
+    let pop_opt = function
+        Empty -> None
+      | Node{l; v; r} -> Some (v, merge l r)
 
     let rec find x = function
         Empty -> raise Not_found

sledge/nonstdlib/ocaml/set.mli

@@ -44,6 +44,8 @@
    of sets of [int * int].
 *)
 
+open! NS0
+
 module type OrderedType =
   sig
     type t
@@ -191,6 +193,11 @@ module type S =
        to the ordering [Ord.compare], where [Ord] is the argument
        given to {!Set.Make}. *)
 
+    val only_elt: t -> elt option
+    (** Return the element of a singleton set, or None otherwise. *)
+
+    val classify : t -> elt zero_one_many
+
     val min_elt: t -> elt
     (** Return the smallest element of the given set
        (with respect to the [Ord.compare] ordering), or raise
@@ -215,16 +222,24 @@ module type S =
 
     val choose: t -> elt
     (** Return one element of the given set, or raise [Not_found] if
-       the set is empty. Which element is chosen is unspecified,
-       but equal elements will be chosen for equal sets. *)
+       the set is empty. Which element is chosen is unspecified, and
+       different elements may be chosen for equal sets. *)
 
     val choose_opt: t -> elt option
     (** Return one element of the given set, or [None] if
-        the set is empty. Which element is chosen is unspecified,
-        but equal elements will be chosen for equal sets.
+        the set is empty. Which element is chosen is unspecified, and
+        different elements may be chosen for equal sets.
         @since 4.05
     *)
 
+    val pop : t -> elt * t
+    (** Find and remove an unspecified element, or raise [Not_found] if the
+        set is empty. [O(1)]. *)
+
+    val pop_opt : t -> (elt * t) option
+    (** Find and remove an unspecified element, or [None] if the set is
+        empty. [O(1)]. *)
+
     val split: elt -> t -> t * bool * t
     (** [split x s] returns a triple [(l, present, r)], where
           [l] is the set of elements of [s] that are