Merge PR #6892: Cleanup implementation of normalize_context_set a bit

2 files changed, 9 insertions, 6 deletions

diff --git a/kernel/uGraph.ml b/kernel/uGraph.ml
index b4ea04a56..e6b27077b 100644
--- a/kernel/uGraph.ml
+++ b/kernel/uGraph.ml
@@ -767,18 +767,18 @@ let normalize_universes g =
     g.entries g
 
 let constraints_of_universes g =
+  let module UF = Unionfind.Make (LSet) (LMap) in
+  let uf = UF.create () in
   let constraints_of u v acc =
     match v with
     | Canonical {univ=u; ltle} ->
       UMap.fold (fun v strict acc->
         let typ = if strict then Lt else Le in
         Constraint.add (u,typ,v) acc) ltle acc
-    | Equiv v -> Constraint.add (u,Eq,v) acc
+    | Equiv v -> UF.union u v uf; acc
   in
-  UMap.fold constraints_of g.entries Constraint.empty
-
-let constraints_of_universes g =
-  constraints_of_universes (normalize_universes g)
+  let csts = UMap.fold constraints_of g.entries Constraint.empty in
+  csts, UF.partition uf
 
 (** [sort_universes g] builds a totally ordered universe graph.  The
     output graph should imply the input graph (and the implication
diff --git a/kernel/uGraph.mli b/kernel/uGraph.mli
index d4fba63fb..cca2eb472 100644
--- a/kernel/uGraph.mli
+++ b/kernel/uGraph.mli
@@ -59,7 +59,10 @@ val empty_universes : t
 
 val sort_universes : t -> t
 
-val constraints_of_universes : t -> Constraint.t
+(** [constraints_of_universes g] returns [csts] and [partition] where
+   [csts] are the non-Eq constraints and [partition] is the partition
+   of the universes into equivalence classes. *)
+val constraints_of_universes : t -> Constraint.t * LSet.t list
 
 val check_subtype : AUContext.t check_function
 (** [check_subtype univ ctx1 ctx2] checks whether [ctx2] is an instance of