From 779e9405c8043cf500ea357cd0f7c1aaf639869b Mon Sep 17 00:00:00 2001
From: Josh Berdine <jjb@fb.com>
Date: Tue, 20 Oct 2020 02:38:06 -0700
Subject: [PATCH] [sledge] Switch from Base.Set to Containers.Set

Summary:
The form of the Base containers interface, in particular the way
comparison functions are passed using Comparators, is slower than
standard functors.

Reviewed By: jvillard

Differential Revision: D24306082

fbshipit-source-id: abf3e0293
---
 sledge/nonstdlib/NS.ml       | 21 +++++++-
 sledge/nonstdlib/NS.mli      | 20 ++++++--
 sledge/nonstdlib/set.ml      | 93 ++++++++++++++++++++++--------------
 sledge/nonstdlib/set_intf.ml | 62 +++++++++++++++++-------
 4 files changed, 138 insertions(+), 58 deletions(-)

diff --git a/sledge/nonstdlib/NS.ml b/sledge/nonstdlib/NS.ml
index e25fc3cc6..39acfe3ee 100644
--- a/sledge/nonstdlib/NS.ml
+++ b/sledge/nonstdlib/NS.ml
@@ -146,10 +146,29 @@ module String = struct
     Core.String :
       sig
         include
-          module type of Core.String with module Map := Core.String.Map
+          module type of Core.String
+            with module Map := Core.String.Map
+            with module Set := Core.String.Set
       end )
 
   module Map = Map.Make (Core.String)
+  module Set = Set.Make (Core.String)
+end
+
+module Int = struct
+  include Stdlib.Int
+
+  include (
+    Int :
+      sig
+        include
+          module type of Core.Int
+            with module Map := Core.Int.Map
+            with module Set := Core.Int.Set
+      end )
+
+  module Map = Map.Make (Int)
+  module Set = Set.Make (Int)
 end
 
 module Q = struct
diff --git a/sledge/nonstdlib/NS.mli b/sledge/nonstdlib/NS.mli
index 39728e93b..67089c549 100644
--- a/sledge/nonstdlib/NS.mli
+++ b/sledge/nonstdlib/NS.mli
@@ -114,11 +114,25 @@ module Qset = Qset
 (** Data types *)
 
 module String : sig
-  include sig
-    include module type of Core.String with module Map := Core.String.Map
-  end
+  include
+    module type of Core.String
+      with module Map := Core.String.Map
+      with module Set := Core.String.Set
 
   module Map : Map.S with type key = string
+  module Set : Set.S with type elt = string
+end
+
+module Int : sig
+  include module type of Stdlib.Int
+
+  include
+    module type of Core.Int
+      with module Map := Core.Int.Map
+      with module Set := Core.Int.Set
+
+  module Map : Map.S with type key = int
+  module Set : Set.S with type elt = int
 end
 
 module Q : sig
diff --git a/sledge/nonstdlib/set.ml b/sledge/nonstdlib/set.ml
index f2a0e0c2e..8dea6aa32 100644
--- a/sledge/nonstdlib/set.ml
+++ b/sledge/nonstdlib/set.ml
@@ -6,57 +6,78 @@
  *)
 
 open NS0
+module Option = CCOpt
 include Set_intf
 
 module Make (Elt : sig
   type t [@@deriving compare, sexp_of]
 end) : S with type elt = Elt.t = struct
-  module EltSet = Core.Set.Make_plain (Elt)
-  module Elt = EltSet.Elt
+  module S = CCSet.Make (Elt)
 
   type elt = Elt.t
+  type t = S.t [@@deriving compare, equal]
 
-  include EltSet.Tree
+  let sexp_of_t s = S.to_list s |> Sexplib.Conv.sexp_of_list Elt.sexp_of_t
 
-  let hash_fold_t hash_fold_elt s m =
-    fold ~f:hash_fold_elt ~init:(Hash.fold_int s (length m)) m
+  module Provide_of_sexp (Elt : sig
+    type t = elt [@@deriving of_sexp]
+  end) =
+  struct
+    let t_of_sexp s =
+      s |> Sexplib.Conv.list_of_sexp Elt.t_of_sexp |> S.of_list
+  end
 
-  let pp ?pre ?suf ?(sep = (",@ " : (unit, unit) fmt)) pp_elt fs x =
-    List.pp ?pre ?suf sep pp_elt fs (elements x)
-
-  let pp_diff pp_elt fs (xs, ys) =
-    let lose = diff xs ys and gain = diff ys xs in
-    if not (is_empty lose) then Format.fprintf fs "-- %a" (pp pp_elt) lose ;
-    if not (is_empty gain) then Format.fprintf fs "++ %a" (pp pp_elt) gain
-
-  let of_ x = add empty x
-  let of_option = Option.fold ~f:add ~init:empty
-  let of_iarray a = of_array (IArray.to_array a)
-  let add_option xo s = Option.fold ~f:add ~init:s xo
-  let add_list xs s = List.fold ~f:add ~init:s xs
+  let empty = S.empty
+  let of_ = S.singleton
+  let of_option xo = Option.map_or S.singleton xo ~default:empty
+  let of_list = S.of_list
+  let add s x = S.add x s
+  let add_option xo s = Option.fold add s xo
+  let add_list xs s = S.add_list s xs
+  let diff = S.diff
+  let inter = S.inter
+  let union = S.union
   let diff_inter s t = (diff s t, inter s t)
+  let union_list ss = List.fold ~f:union ~init:empty ss
+  let is_empty = S.is_empty
+  let mem s x = S.mem x s
+  let is_subset s ~of_:t = S.subset s t
+  let disjoint = S.disjoint
+  let max_elt = S.max_elt_opt
 
-  let rec disjoint s1 s2 =
-    match choose s1 with
-    | None -> true
-    | _ when is_empty s2 -> true
-    | _ when s1 == s2 -> false
-    | Some x -> (
-        let l1, _, r1 = split s1 x in
-        match split s2 x with
-        | _, Some _, _ -> false
-        | l2, None, r2 -> disjoint l1 l2 && disjoint r1 r2 )
-
-  let choose_exn s =
-    let@ {return} = with_return in
-    binary_search_segmented s `Last_on_left ~segment_of:return |> ignore ;
-    raise (Not_found_s (Atom __LOC__))
+  let root_elt s =
+    let exception Found in
+    let found = ref None in
+    try
+      S.for_all
+        (fun elt ->
+          found := Some elt ;
+          raise Found )
+        s
+      |> ignore ;
+      None
+    with Found -> !found
 
-  let choose s = try Some (choose_exn s) with Not_found_s _ -> None
+  let choose = root_elt
+  let choose_exn m = Option.get_exn (choose m)
 
   let pop_exn s =
     let elt = choose_exn s in
-    (elt, remove s elt)
+    (elt, S.remove elt s)
 
-  let pop s = choose s |> Option.map ~f:(fun elt -> (elt, remove s elt))
+  let elements = S.elements
+  let map s ~f = S.map f s
+  let filter s ~f = S.filter f s
+  let iter s ~f = S.iter f s
+  let exists s ~f = S.exists f s
+  let for_all s ~f = S.for_all f s
+  let fold s ~init ~f = S.fold (fun x a -> f a x) s init
+
+  let pp ?pre ?suf ?(sep = (",@ " : (unit, unit) fmt)) pp_elt fs x =
+    List.pp ?pre ?suf sep pp_elt fs (S.elements x)
+
+  let pp_diff pp_elt fs (xs, ys) =
+    let lose = diff xs ys and gain = diff ys xs in
+    if not (is_empty lose) then Format.fprintf fs "-- %a" (pp pp_elt) lose ;
+    if not (is_empty gain) then Format.fprintf fs "++ %a" (pp pp_elt) gain
 end
diff --git a/sledge/nonstdlib/set_intf.ml b/sledge/nonstdlib/set_intf.ml
index 0bdca1ed1..a62817542 100644
--- a/sledge/nonstdlib/set_intf.ml
+++ b/sledge/nonstdlib/set_intf.ml
@@ -9,36 +9,62 @@ open NS0
 
 module type S = sig
   type elt
+  type t [@@deriving compare, equal, sexp_of]
 
-  module Elt : sig
-    type t = elt
-
-    include Comparator.S with type t := t
+  module Provide_of_sexp (_ : sig
+    type t = elt [@@deriving of_sexp]
+  end) : sig
+    val t_of_sexp : Sexp.t -> t
   end
 
-  include Core_kernel.Set_intf.Make_S_plain_tree(Elt).S
-
-  val hash_fold_t : elt Hash.folder -> t Hash.folder
-
-  val pp :
-       ?pre:(unit, unit) fmt
-    -> ?suf:(unit, unit) fmt
-    -> ?sep:(unit, unit) fmt
-    -> elt pp
-    -> t pp
+  (** {1 Construct} *)
 
-  val pp_diff : elt pp -> (t * t) pp
+  val empty : t
   val of_ : elt -> t
   val of_option : elt option -> t
-  val of_iarray : elt IArray.t -> t
+  val of_list : elt list -> t
+  val add : t -> elt -> t
   val add_option : elt option -> t -> t
   val add_list : elt list -> t -> t
+  val diff : t -> t -> t
+  val inter : t -> t -> t
+  val union : t -> t -> t
   val diff_inter : t -> t -> t * t
+  val union_list : t list -> t
+
+  (** {1 Query} *)
+
+  val is_empty : t -> bool
+  val mem : t -> elt -> bool
+  val is_subset : t -> of_:t -> bool
   val disjoint : t -> t -> bool
+  val max_elt : t -> elt option
 
   val pop_exn : t -> elt * t
   (** Find and remove an unspecified element. [O(1)]. *)
 
-  val pop : t -> (elt * t) option
-  (** Find and remove an unspecified element. [O(1)]. *)
+  val elements : t -> elt list
+
+  (** {1 Transform} *)
+
+  val map : t -> f:(elt -> elt) -> t
+  val filter : t -> f:(elt -> bool) -> t
+
+  (** {1 Traverse} *)
+
+  val iter : t -> f:(elt -> unit) -> unit
+  val exists : t -> f:(elt -> bool) -> bool
+  val for_all : t -> f:(elt -> bool) -> bool
+  val fold : t -> init:'a -> f:('a -> elt -> 'a) -> 'a
+
+  (** {1 Pretty-print} *)
+
+  val pp :
+       ?pre:(unit, unit) fmt
+    -> ?suf:(unit, unit) fmt
+    -> ?sep:(unit, unit) fmt
+    -> elt pp
+    -> t pp
+
+  val pp_diff : elt pp -> (t * t) pp
 end