Union-Find: function to extract the list of partitions

src/stages/4-ast_typed/PP_generic.ml

@@ -87,8 +87,12 @@ let op ppf = {
       | None -> fprintf ppf "None"
       | Some v -> fprintf ppf "%a" (fun _ppf -> continue ()) v) ;
     poly_unionfind            = (fun _visitor continue () p   ->
-      let lst = (UnionFind.Poly2.elements p) in
+      let lst = (UnionFind.Poly2.partitions p) in
-      fprintf ppf "LMap [ %a ]" (list_sep (fun _ppf -> continue ()) (fun ppf () -> fprintf ppf " ; ")) lst);
+      let aux1 l = fprintf ppf "[@,@[<hv 2> (*%a*) %a @]@,]"
+                     (fun _ppf -> continue ()) (UnionFind.Poly2.repr (List.hd l) p)
+                     (list_sep (fun _ppf -> continue ()) (fun ppf () -> fprintf ppf " ;@ ")) l in
+      let aux2 = list_sep (fun _ppf -> aux1) (fun ppf () -> fprintf ppf " ;@ ") in
+      fprintf ppf "UnionFind [@,@[<hv 2> %a @]@,]" aux2 lst);
     poly_set                  = (fun _visitor continue () set   ->
       let lst = (RedBlackTrees.PolySet.elements set) in
       fprintf ppf "Set [@,@[<hv 2> %a @]@,]" (list_sep (fun _ppf -> continue ()) (fun ppf () -> fprintf ppf " ;@ ")) lst);

vendors/UnionFind/Poly2.ml

@@ -1,8 +1,6 @@
 (** Persistent implementation of the Union/Find algorithm with
     height-balanced forests and no path compression. *)
 
-(* type item = Item.t *)
-
 let equal compare i j = compare i j = 0
 
 type height = int
@@ -123,6 +121,30 @@ let elements : 'item . 'item partition -> 'item list =
   fun { to_string=_; compare=_; map } ->
   map_sorted_keys map
 
+let partitions : 'item . 'item partition -> 'item list list =
+  let compare_lists_by_first cmp la lb =
+    match la,lb with
+      | [],[] -> 0
+      | [],_ -> -1
+      | _,[] -> 1
+      | a::_, b::_ -> cmp a b in
+  fun ({ to_string=_; compare; map } as p) ->
+  let aux acc elt =
+    RedBlackTrees.PolyMap.update
+      (repr elt p)
+      (function None -> Some [elt] | Some l -> Some (elt::l))
+      acc in
+  let grouped = List.fold_left
+    aux
+    (RedBlackTrees.PolyMap.create ~cmp:compare)
+    (map_sorted_keys map) in
+  let partitions = RedBlackTrees.PolyMap.bindings grouped in
+  (* Sort the elements within partitions and partitions by their smallest element *)
+  let partitions = List.map snd partitions in
+  let partitions = List.map (List.sort compare) partitions in
+  let partitions = List.sort (compare_lists_by_first compare) partitions in
+  partitions
+
 (** {1 Printing} *)
 
 let print ppf (p: 'item partition) =

vendors/UnionFind/Poly2.mli

@@ -47,6 +47,13 @@ val mem : 'item -> 'item partition -> 'item option
     ordered in ascending order *)
 val elements : 'item partition -> 'item list
 
+(** The value of the call [partitions p] is a list of the partitions
+   of p. A partition is a list of elements. The elements and
+   partitions are returned with a deterministic order (regardless of
+   the way the aliases have been made, the same partition will always
+   have the same order). *)
+val partitions : 'item partition -> 'item list list
+
 (** The call [print p] is a value of type [Buffer.t] containing
         strings denoting the partition [p], based on
         [Ord.to_string]. *)