Reviewed By: ddino Differential Revision: D8348274 fbshipit-source-id: 6bd6ce6master
Mehdi Bouaziz
7 years ago
committed by
Facebook Github Bot
parent
21ced6af62
commit
9ae3b42aef
@ -0,0 +1,149 @@
|
||||
(*
|
||||
* Copyright (c) 2018-present, Facebook, Inc.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*)
|
||||
|
||||
open! IStd
|
||||
|
||||
module type Set = sig
|
||||
type elt [@@deriving compare]
|
||||
|
||||
type t
|
||||
|
||||
val create : elt -> t
|
||||
|
||||
val compare_size : t -> t -> int
|
||||
|
||||
val merge : from:t -> to_:t -> unit
|
||||
end
|
||||
|
||||
module Make (Set : Set) = struct
|
||||
module H = struct
|
||||
include Caml.Hashtbl
|
||||
|
||||
let fold : (('a, 'b) t, 'a * 'b, 'accum) Container.fold =
|
||||
fun h ~init ~f ->
|
||||
let f' k v accum = f accum (k, v) in
|
||||
fold f' h init
|
||||
end
|
||||
|
||||
module Repr : sig
|
||||
(* Sort-of abstracting away the fact that a representative is just an element itself.
|
||||
This ensures that the [Sets] hashtable is accessed with representative only. *)
|
||||
|
||||
type t = private Set.elt
|
||||
|
||||
val equal : t -> t -> bool
|
||||
|
||||
val of_elt : Set.elt -> t
|
||||
|
||||
val is_simpler_than : t -> t -> bool
|
||||
end = struct
|
||||
type t = Set.elt [@@deriving compare]
|
||||
|
||||
let equal = [%compare.equal : t]
|
||||
|
||||
let of_elt e = e
|
||||
|
||||
let is_simpler_than r1 r2 = compare r1 r2 <= 0
|
||||
end
|
||||
|
||||
module Reprs = struct
|
||||
type t = (Set.elt, Repr.t) H.t
|
||||
|
||||
let create () = H.create 1
|
||||
|
||||
let is_a_repr (t: t) e = not (H.mem t e)
|
||||
|
||||
let rec find (t: t) e : Repr.t =
|
||||
match H.find_opt t e with
|
||||
| None ->
|
||||
Repr.of_elt e
|
||||
| Some r ->
|
||||
let r' = find t (r :> Set.elt) in
|
||||
if not (phys_equal r r') then H.replace t e r' ;
|
||||
r'
|
||||
|
||||
|
||||
let merge (t: t) ~(from: Repr.t) ~(to_: Repr.t) = H.replace t (from :> Set.elt) to_
|
||||
end
|
||||
|
||||
module Sets = struct
|
||||
type t = (Repr.t, Set.t) H.t
|
||||
|
||||
let create () = H.create 1
|
||||
|
||||
let find_create t (r: Repr.t) =
|
||||
match H.find_opt t r with
|
||||
| Some set ->
|
||||
set
|
||||
| None ->
|
||||
let set = Set.create (r :> Set.elt) in
|
||||
H.replace t r set ; set
|
||||
|
||||
|
||||
let fold = H.fold
|
||||
|
||||
let remove_now t r = H.remove t r
|
||||
end
|
||||
|
||||
(**
|
||||
Data-structure for disjoint sets.
|
||||
[reprs] is the mapping element -> representative
|
||||
[sets] is the mapping representative -> set
|
||||
|
||||
It implements path-compression and union by size, hence find and union are amortized O(1)-ish.
|
||||
|
||||
[nb_iterators] and [to_remove] are used to defer removing elements to avoid iterator invalidation during fold.
|
||||
*)
|
||||
type t = {reprs: Reprs.t; sets: Sets.t; mutable nb_iterators: int; mutable to_remove: Repr.t list}
|
||||
|
||||
let create () = {reprs= Reprs.create (); sets= Sets.create (); nb_iterators= 0; to_remove= []}
|
||||
|
||||
let find t e = Reprs.find t.reprs e
|
||||
|
||||
let do_merge t ~from_r ~from_set ~to_r ~to_set =
|
||||
Reprs.merge t.reprs ~from:from_r ~to_:to_r ;
|
||||
Set.merge ~from:from_set ~to_:to_set ;
|
||||
if t.nb_iterators <= 0 then Sets.remove_now t.sets from_r
|
||||
else t.to_remove <- from_r :: t.to_remove
|
||||
|
||||
|
||||
let union t e1 e2 =
|
||||
let repr1 = find t e1 in
|
||||
let repr2 = find t e2 in
|
||||
if Repr.equal repr1 repr2 then None
|
||||
else
|
||||
let set1 = Sets.find_create t.sets repr1 in
|
||||
let set2 = Sets.find_create t.sets repr2 in
|
||||
let cmp_size = Set.compare_size set1 set2 in
|
||||
if cmp_size < 0 || (Int.equal cmp_size 0 && Repr.is_simpler_than repr2 repr1) then (
|
||||
(* A desired side-effect of using [is_simpler_than] is that the representative for a set will always be a [`Node]. For now. *)
|
||||
do_merge t ~from_r:repr1 ~from_set:set1 ~to_r:repr2 ~to_set:set2 ;
|
||||
Some (e1, e2) )
|
||||
else (
|
||||
do_merge t ~from_r:repr2 ~from_set:set2 ~to_r:repr1 ~to_set:set1 ;
|
||||
Some (e2, e1) )
|
||||
|
||||
|
||||
let is_still_a_repr t ((repr: Repr.t), _) = Reprs.is_a_repr t.reprs (repr :> Set.elt)
|
||||
|
||||
let after_fold t =
|
||||
let new_nb_iterators = t.nb_iterators - 1 in
|
||||
t.nb_iterators <- new_nb_iterators ;
|
||||
if new_nb_iterators <= 0 && not (List.is_empty t.to_remove) then (
|
||||
List.iter t.to_remove ~f:(Sets.remove_now t.sets) ;
|
||||
t.to_remove <- [] )
|
||||
|
||||
|
||||
let fold_sets t ~init ~f =
|
||||
t.nb_iterators <- t.nb_iterators + 1 ;
|
||||
match IContainer.filter ~fold:Sets.fold ~filter:(is_still_a_repr t) t.sets ~init ~f with
|
||||
| result ->
|
||||
after_fold t ; result
|
||||
| exception e ->
|
||||
(* Ensures [nb_iterators] is correct *)
|
||||
IExn.reraise_after ~f:(fun () -> after_fold t) e
|
||||
end
|
||||
@ -0,0 +1,40 @@
|
||||
(*
|
||||
* Copyright (c) 2018-present, Facebook, Inc.
|
||||
*
|
||||
* This source code is licensed under the MIT license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*)
|
||||
|
||||
open! IStd
|
||||
|
||||
module type Set = sig
|
||||
type elt [@@deriving compare]
|
||||
|
||||
type t
|
||||
|
||||
val create : elt -> t
|
||||
|
||||
val compare_size : t -> t -> int
|
||||
|
||||
val merge : from:t -> to_:t -> unit
|
||||
end
|
||||
|
||||
module Make (Set : Set) : sig
|
||||
module Repr : sig
|
||||
type t = private Set.elt
|
||||
|
||||
val equal : t -> t -> bool
|
||||
end
|
||||
|
||||
type t
|
||||
|
||||
val create : unit -> t
|
||||
|
||||
val find : t -> Set.elt -> Repr.t
|
||||
|
||||
val union : t -> Set.elt -> Set.elt -> (Set.elt * Set.elt) option
|
||||
(** [union t e1 e2] returns [None] if [e1] and [e2] were already in the same set, [Some (a, b)] if [a] is merged into [b] (were [(a, b)] is either [(e1, e2)] or [(e2, e1)]). *)
|
||||
|
||||
val fold_sets : (t, Repr.t * Set.t, 'accum) Container.fold
|
||||
(** It is safe to call [find] or [union] while folding. *)
|
||||
end
|
||||
Loading…
Reference in new issue