Merge PR #13718: Move printing and sorting out of AcyclicGraph

Reviewed-by: SkySkimmer

7 files changed, 120 insertions, 161 deletions

diff --git a/checker/checker.ml b/checker/checker.ml
index 08d92bb7b3..bdfc5f07be 100644
--- a/checker/checker.ml
+++ b/checker/checker.ml
@@ -289,7 +289,7 @@ let explain_exn = function
                              Constr.debug_print a ++ fnl ());
         Feedback.msg_notice (str"====== universes ====" ++ fnl () ++
                              (UGraph.pr_universes Univ.Level.pr
-                                (ctx.Environ.env_stratification.Environ.env_universes)));
+                                (UGraph.repr (ctx.Environ.env_stratification.Environ.env_universes))));
         str "CantApplyBadType at argument " ++ int n
       | CantApplyNonFunctional _ -> str"CantApplyNonFunctional"
       | IllFormedRecBody _ -> str"IllFormedRecBody"
diff --git a/dev/top_printers.ml b/dev/top_printers.ml
index c4eed5756f..f3d6239c6f 100644
--- a/dev/top_printers.ml
+++ b/dev/top_printers.ml
@@ -240,7 +240,7 @@ let ppuniverseconstraints c = pp (UnivProblem.Set.pr c)
 let ppuniverse_context_future c =
   let ctx = Future.force c in
     ppuniverse_context ctx
-let ppuniverses u = pp (UGraph.pr_universes Level.pr u)
+let ppuniverses u = pp (UGraph.pr_universes Level.pr (UGraph.repr u))
 let ppnamedcontextval e =
   let env = Global.env () in
   let sigma = Evd.from_env env in
diff --git a/kernel/uGraph.ml b/kernel/uGraph.ml
index 096e458ec4..b988ec40a7 100644
--- a/kernel/uGraph.ml
+++ b/kernel/uGraph.ml
@@ -222,15 +222,35 @@ let choose p g u = if Level.is_sprop u
   then if p u then Some u else None
   else G.choose p g.graph u
 
-let dump_universes f g = G.dump f g.graph
-
 let check_universes_invariants g = G.check_invariants ~required_canonical:Level.is_small g.graph
 
-let pr_universes prl g = G.pr prl g.graph
-
-let dummy_mp = Names.DirPath.make [Names.Id.of_string "Type"]
-let make_dummy i = Level.(make (UGlobal.make dummy_mp i))
-let sort_universes g = g_map (G.sort make_dummy [Level.prop;Level.set]) g
+(** Pretty-printing *)
+
+let pr_pmap sep pr map =
+  let cmp (u,_) (v,_) = Level.compare u v in
+  Pp.prlist_with_sep sep pr (List.sort cmp (LMap.bindings map))
+
+let pr_arc prl = let open Pp in
+  function
+  | u, G.Node ltle ->
+    if LMap.is_empty ltle then mt ()
+    else
+      prl u ++ str " " ++
+      v 0
+        (pr_pmap spc (fun (v, strict) ->
+              (if strict then str "< " else str "<= ") ++ prl v)
+            ltle) ++
+      fnl ()
+  | u, G.Alias v ->
+    prl u  ++ str " = " ++ prl v ++ fnl ()
+
+type node = G.node =
+| Alias of Level.t
+| Node of bool LMap.t
+
+let repr g = G.repr g.graph
+
+let pr_universes prl g = pr_pmap Pp.mt (pr_arc prl) g
 
 (** Profiling *)
 
diff --git a/kernel/uGraph.mli b/kernel/uGraph.mli
index 87b3634e28..9ac29f5139 100644
--- a/kernel/uGraph.mli
+++ b/kernel/uGraph.mli
@@ -77,15 +77,9 @@ exception UndeclaredLevel of Univ.Level.t
 
 val check_declared_universes : t -> Univ.LSet.t -> unit
 
-(** {6 Pretty-printing of universes. } *)
-
-val pr_universes : (Level.t -> Pp.t) -> t -> Pp.t
-
 (** The empty graph of universes *)
 val empty_universes : t
 
-val sort_universes : t -> 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. *)
@@ -108,10 +102,17 @@ val check_subtype : lbound:Bound.t -> AUContext.t check_function
 (** [check_subtype univ ctx1 ctx2] checks whether [ctx2] is an instance of
     [ctx1]. *)
 
-(** {6 Dumping to a file } *)
+(** {6 Dumping} *)
+
+type node =
+| Alias of Level.t
+| Node of bool LMap.t (** Nodes v s.t. u < v (true) or u <= v (false) *)
+
+val repr : t -> node LMap.t
+
+(** {6 Pretty-printing of universes. } *)
 
-val dump_universes :
-  (constraint_type -> Level.t -> Level.t -> unit) -> t -> unit
+val pr_universes : (Level.t -> Pp.t) -> node LMap.t -> Pp.t
 
 (** {6 Debugging} *)
 val check_universes_invariants : t -> unit
diff --git a/lib/acyclicGraph.ml b/lib/acyclicGraph.ml
index 8da09dc98a..14c08da35d 100644
--- a/lib/acyclicGraph.ml
+++ b/lib/acyclicGraph.ml
@@ -76,8 +76,6 @@ module Make (Point:Point) = struct
       mutable status: status
     }
 
-  let big_rank = 1000000
-
   (* A Point.t is either an alias for another one, or a canonical one,
      for which we know the points that are above *)
 
@@ -158,30 +156,6 @@ module Make (Point:Point) = struct
     assert (g.index > min_int);
     { g with index = g.index - 1 }
 
-  (* [safe_repr] is like [repr] but if the graph doesn't contain the
-     searched point, we add it. *)
-  let safe_repr g u =
-    let rec safe_repr_rec entries u =
-      match PMap.find u entries with
-      | Equiv v -> safe_repr_rec entries v
-      | Canonical arc -> arc
-    in
-    try g, safe_repr_rec g.entries u
-    with Not_found ->
-      let can =
-        { canon = u;
-          ltle = PMap.empty; gtge = PSet.empty;
-          rank = 0;
-          klvl = 0; ilvl = 0;
-          status = NoMark }
-      in
-      let g = { g with
-                entries = PMap.add u (Canonical can) g.entries;
-                n_nodes = g.n_nodes + 1 }
-      in
-      let g = use_index g u in
-      g, repr g u
-
   (* Returns 1 if u is higher than v in topological order.
              -1        lower
              0 if u = v *)
@@ -676,29 +650,6 @@ module Make (Point:Point) = struct
 
   (* Normalization *)
 
-  (** [normalize g] returns a graph where all edges point
-      directly to the canonical representent of their target. The output
-      graph should be equivalent to the input graph from a logical point
-      of view, but optimized. We maintain the invariant that the key of
-      a [Canonical] element is its own name, by keeping [Equiv] edges. *)
-  let normalize g =
-    let g =
-      { g with
-        entries = PMap.map (fun entry ->
-            match entry with
-            | Equiv u -> Equiv ((repr g u).canon)
-            | Canonical ucan -> Canonical { ucan with rank = 1 })
-            g.entries }
-    in
-    PMap.fold (fun _ u g ->
-        match u with
-        | Equiv _u -> g
-        | Canonical u ->
-          let _, u, g = get_ltle g u in
-          let _, _, g = get_gtge g u in
-          g)
-      g.entries g
-
   let constraints_of g =
     let module UF = Unionfind.Make (PSet) (PMap) in
     let uf = UF.create () in
@@ -769,85 +720,14 @@ module Make (Point:Point) = struct
         ) g.entries; None
       with Found v -> Some v
 
-  let sort make_dummy first g =
-    let cans =
-      PMap.fold (fun _ u l ->
-          match u with
-          | Equiv _ -> l
-          | Canonical can -> can :: l
-        ) g.entries []
-    in
-    let cans = List.sort topo_compare cans in
-    let lowest =
-      PMap.mapi (fun u _ -> if CList.mem_f Point.equal u first then 0 else 2)
-        (PMap.filter
-           (fun _ u -> match u with Equiv _ -> false | Canonical _ -> true)
-           g.entries)
-    in
-    let lowest =
-      List.fold_left (fun lowest can ->
-          let lvl = PMap.find can.canon lowest in
-          PMap.fold (fun u' strict lowest ->
-              let cost = if strict then 1 else 0 in
-              let u' = (repr g u').canon in
-              PMap.modify u' (fun _ lvl0 -> max lvl0 (lvl+cost)) lowest)
-            can.ltle lowest)
-        lowest cans
-    in
-    let max_lvl = PMap.fold (fun _ a b -> max a b) lowest 0 in
-    let types = Array.init (max_lvl + 1) (fun i ->
-        match List.nth_opt first i with
-        | Some u -> u
-        | None -> make_dummy (i-2))
-    in
-    let g = Array.fold_left (fun g u ->
-        let g, u = safe_repr g u in
-        change_node g { u with rank = big_rank }) g types
-    in
-    let g = if max_lvl > List.length first && not (CList.is_empty first) then
-        enforce_lt (CList.last first) types.(List.length first) g
-      else g
-    in
-    let g =
-      PMap.fold (fun u lvl g -> enforce_eq u (types.(lvl)) g)
-        lowest g
-    in
-    normalize g
-
-  (** Pretty-printing *)
-
-  let pr_pmap sep pr map =
-    let cmp (u,_) (v,_) = Point.compare u v in
-    Pp.prlist_with_sep sep pr (List.sort cmp (PMap.bindings map))
+  type node = Alias of Point.t | Node of bool Point.Map.t
+  type repr = node Point.Map.t
 
-  let pr_arc prl = let open Pp in
-    function
-    | _, Canonical {canon=u; ltle; _} ->
-      if PMap.is_empty ltle then mt ()
-      else
-        prl u ++ str " " ++
-        v 0
-          (pr_pmap spc (fun (v, strict) ->
-               (if strict then str "< " else str "<= ") ++ prl v)
-              ltle) ++
-        fnl ()
-    | u, Equiv v ->
-      prl u  ++ str " = " ++ prl v ++ fnl ()
-
-  let pr prl g =
-    pr_pmap Pp.mt (pr_arc prl) g.entries
-
-  (* Dumping constraints to a file *)
-
-  let dump output g =
-    let dump_arc u = function
-      | Canonical {canon=u; ltle; _} ->
-        PMap.iter (fun v strict ->
-            let typ = if strict then Lt else Le in
-            output typ u v) ltle;
-      | Equiv v ->
-        output Eq u v
+  let repr g =
+    let map n = match n with
+    | Canonical n -> Node n.ltle
+    | Equiv u -> Alias u
     in
-    PMap.iter dump_arc g.entries
+    Point.Map.map map g.entries
 
 end
diff --git a/lib/acyclicGraph.mli b/lib/acyclicGraph.mli
index e9f05ed74d..8c9d2e6461 100644
--- a/lib/acyclicGraph.mli
+++ b/lib/acyclicGraph.mli
@@ -65,18 +65,12 @@ module Make (Point:Point) : sig
 
   val choose : (Point.t -> bool) -> t -> Point.t -> Point.t option
 
-  val sort : (int -> Point.t) -> Point.t list -> t -> t
-  (** [sort mk first g] builds a totally ordered graph. The output
-     graph should imply the input graph (and the implication will be
-     strict most of the time), but is not necessarily minimal. The
-     lowest points in the result are identified with [first].
-     Moreover, it adds levels [Type.n] to identify the points (not in
-     [first]) at level n. An artificial constraint (last first < mk
-     (length first)) is added to ensure that they are not merged.
-     Note: the result is unspecified if the input graph already
-     contains [mk n] nodes. *)
-
-  val pr : (Point.t -> Pp.t) -> t -> Pp.t
-
-  val dump : (constraint_type -> Point.t -> Point.t -> unit) -> t -> unit
+  (** {5 High-level representation} *)
+
+  type node =
+  | Alias of Point.t
+  | Node of bool Point.Map.t (** Nodes v s.t. u < v (true) or u <= v (false) *)
+  type repr = node Point.Map.t
+  val repr : t -> repr
+
 end
diff --git a/vernac/vernacentries.ml b/vernac/vernacentries.ml
index e8cb1d65a9..4f3fc46c12 100644
--- a/vernac/vernacentries.ml
+++ b/vernac/vernacentries.ml
@@ -309,6 +309,17 @@ let print_registered () =
   in
   hov 0 (prlist_with_sep fnl pr_lib_ref @@ Coqlib.get_lib_refs ())
 
+let dump_universes output g =
+  let open Univ in
+  let dump_arc u = function
+    | UGraph.Node ltle ->
+      Univ.LMap.iter (fun v strict ->
+          let typ = if strict then Lt else Le in
+          output typ u v) ltle;
+    | UGraph.Alias v ->
+      output Eq u v
+  in
+  Univ.LMap.iter dump_arc g
 
 let dump_universes_gen prl g s =
   let output = open_out s in
@@ -342,7 +353,7 @@ let dump_universes_gen prl g s =
   in
   let output_constraint k l r = output_constraint k (prl l) (prl r) in
   try
-    UGraph.dump_universes output_constraint g;
+    dump_universes output_constraint g;
     close ();
     str "Universes written to file \"" ++ str s ++ str "\"."
   with reraise ->
@@ -367,13 +378,66 @@ let universe_subgraph ?loc kept univ =
   let univ = LSet.fold add kept UGraph.initial_universes in
   UGraph.merge_constraints csts univ
 
+let sort_universes g =
+  let open Univ in
+  let rec normalize u = match LMap.find u g with
+  | UGraph.Alias u -> normalize u
+  | UGraph.Node _ -> u
+  in
+  let get_next u = match LMap.find u g with
+  | UGraph.Alias u -> assert false (* nodes are normalized *)
+  | UGraph.Node ltle -> ltle
+  in
+  (* Compute the longest chain of Lt constraints from Set to any universe *)
+  let rec traverse accu todo = match todo with
+  | [] -> accu
+  | (u, n) :: todo ->
+    let () = assert (Level.equal (normalize u) u) in
+    let n = match LMap.find u accu with
+    | m -> if m < n then Some n else None
+    | exception Not_found -> Some n
+    in
+    match n with
+    | None -> traverse accu todo
+    | Some n ->
+      let accu = LMap.add u n accu in
+      let next = get_next u in
+      let fold v lt todo =
+        let v = normalize v in
+        if lt then (v, n + 1) :: todo else (v, n) :: todo
+      in
+      let todo = LMap.fold fold next todo in
+      traverse accu todo
+  in
+  (* Only contains normalized nodes *)
+  let levels = traverse LMap.empty [normalize Level.set, 0] in
+  let max_level = LMap.fold (fun _ n accu -> max n accu) levels 0 in
+  let dummy_mp = Names.DirPath.make [Names.Id.of_string "Type"] in
+  let ulevels = Array.init max_level (fun i -> Level.(make (UGlobal.make dummy_mp i))) in
+  let ulevels = Array.cons Level.set ulevels in
+  (* Add the normal universes *)
+  let fold (cur, ans) u =
+    let ans = LMap.add cur (UGraph.Node (LMap.singleton u true)) ans in
+    (u, ans)
+  in
+  let _, ans = Array.fold_left fold (Level.prop, LMap.empty) ulevels in
+  (* Add alias pointers *)
+  let fold u _ ans =
+    if Level.is_small u then ans
+    else
+      let n = LMap.find (normalize u) levels in
+      LMap.add u (UGraph.Alias ulevels.(n)) ans
+  in
+  LMap.fold fold g ans
+
 let print_universes ?loc ~sort ~subgraph dst =
   let univ = Global.universes () in
   let univ = match subgraph with
     | None -> univ
     | Some g -> universe_subgraph ?loc g univ
   in
-  let univ = if sort then UGraph.sort_universes univ else univ in
+  let univ = UGraph.repr univ in
+  let univ = if sort then sort_universes univ else univ in
   let pr_remaining =
     if Global.is_joined_environment () then mt ()
     else str"There may remain asynchronous universe constraints"