Universe level maps use HMaps.

2 files changed, 48 insertions, 4 deletions

diff --git a/lib/unionfind.ml b/lib/unionfind.ml
index 38ab0ffa6a..5c9bb3af9e 100644
--- a/lib/unionfind.ml
+++ b/lib/unionfind.ml
@@ -53,7 +53,28 @@ module type PartitionSig = sig
 
 end
 
-module Make (S:Set.S)(M:Map.S with type key = S.elt) = struct
+module type SetS =
+sig
+  type t
+  type elt
+  val singleton : elt -> t
+  val union : t -> t -> t
+  val choose : t -> elt
+  val iter : (elt -> unit) -> t -> unit
+end
+
+module type MapS =
+sig
+  type key
+  type +'a t
+  val empty : 'a t
+  val find : key -> 'a t -> 'a
+  val add : key -> 'a -> 'a t -> 'a t
+  val mem : key -> 'a t -> bool
+  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
+end
+
+module Make (S:SetS)(M:MapS with type key = S.elt) = struct
 
   type elt = S.elt
   type set = S.t
@@ -101,7 +122,7 @@ module Make (S:Set.S)(M:Map.S with type key = S.elt) = struct
 
   let union_set s p =
     try
-      let x = S.min_elt s in
+      let x = S.choose s in
       S.iter (fun y -> union x y p) s
     with Not_found -> ()
 
diff --git a/lib/unionfind.mli b/lib/unionfind.mli
index 18468661ae..cdd483e4ca 100644
--- a/lib/unionfind.mli
+++ b/lib/unionfind.mli
@@ -51,7 +51,30 @@ module type PartitionSig = sig
 
 end
 
+module type SetS =
+sig
+  type t
+  type elt
+  val singleton : elt -> t
+  val union : t -> t -> t
+  val choose : t -> elt
+  val iter : (elt -> unit) -> t -> unit
+end
+(** Minimal interface for sets, subtype of stdlib's Set. *)
+
+module type MapS =
+sig
+  type key
+  type +'a t
+  val empty : 'a t
+  val find : key -> 'a t -> 'a
+  val add : key -> 'a -> 'a t -> 'a t
+  val mem : key -> 'a t -> bool
+  val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
+end
+(** Minimal interface for maps, subtype of stdlib's Map. *)
+
 module Make :
-  functor (S:Set.S) ->
-    functor (M:Map.S with type key = S.elt) ->
+  functor (S:SetS) ->
+    functor (M:MapS with type key = S.elt) ->
       PartitionSig with type elt = S.elt and type set = S.t