From d298eb1bada97c91db83e70d9751940dcb339f30 Mon Sep 17 00:00:00 2001 From: Josh Berdine Date: Sun, 21 Feb 2021 13:18:41 -0800 Subject: [PATCH] [sledge] Optimize Set operations Reviewed By: jvillard Differential Revision: D26338013 fbshipit-source-id: 3a21efa9b --- sledge/nonstdlib/NSSet.ml | 48 +++++----------------------------- sledge/nonstdlib/NSSet_intf.ml | 8 ++++++ sledge/nonstdlib/ocaml/set.ml | 30 +++++++++++++++++++-- sledge/nonstdlib/ocaml/set.mli | 23 +++++++++++++--- 4 files changed, 61 insertions(+), 48 deletions(-) diff --git a/sledge/nonstdlib/NSSet.ml b/sledge/nonstdlib/NSSet.ml index ca171fabf..4365913cb 100644 --- a/sledge/nonstdlib/NSSet.ml +++ b/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 diff --git a/sledge/nonstdlib/NSSet_intf.ml b/sledge/nonstdlib/NSSet_intf.ml index ca26c3507..03e33ac98 100644 --- a/sledge/nonstdlib/NSSet_intf.ml +++ b/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)]. *) diff --git a/sledge/nonstdlib/ocaml/set.ml b/sledge/nonstdlib/ocaml/set.ml index 036ae264e..d22641cf5 100644 --- a/sledge/nonstdlib/ocaml/set.ml +++ b/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 diff --git a/sledge/nonstdlib/ocaml/set.mli b/sledge/nonstdlib/ocaml/set.mli index 669a344fe..949a3e27c 100644 --- a/sledge/nonstdlib/ocaml/set.mli +++ b/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