More efficient implementation of equality-up-to-universes in Universes.

Instead of accumulating constraints which are not present in the original
graph, we parametrize the equality function by a function actually merging
those constraints in the current graph. This prevents doing the work twice.

5 files changed, 55 insertions, 51 deletions

diff --git a/engine/evd.ml b/engine/evd.ml
index 069fcbfa6..00a869fda 100644
--- a/engine/evd.ml
+++ b/engine/evd.ml
@@ -962,11 +962,13 @@ let test_conversion env d pb t u =
 exception UniversesDiffer = UState.UniversesDiffer
 
 let eq_constr_univs evd t u =
-  let b, c = Universes.eq_constr_univs_infer (UState.ugraph evd.universes) t u in
-    if b then 
-      try let evd' = add_universe_constraints evd c in evd', b
-      with Univ.UniverseInconsistency _ | UniversesDiffer -> evd, false
-    else evd, b
+  let fold cstr sigma =
+    try Some (add_universe_constraints sigma cstr)
+    with Univ.UniverseInconsistency _ | UniversesDiffer -> None
+  in
+  match Universes.eq_constr_univs_infer (UState.ugraph evd.universes) fold t u evd with
+  | None -> evd, false
+  | Some evd -> evd, true
 
 let e_eq_constr_univs evdref t u =
   let evd, b = eq_constr_univs !evdref t u in
diff --git a/library/universes.ml b/library/universes.ml
index 225e65842..a157a747c 100644
--- a/library/universes.ml
+++ b/library/universes.ml
@@ -102,6 +102,7 @@ module Constraints = struct
 end
 
 type universe_constraints = Constraints.t
+type 'a constraint_accumulator = universe_constraints -> 'a -> 'a option
 type 'a universe_constrained = 'a * universe_constraints
 
 type 'a universe_constraint_function = 'a -> 'a -> universe_constraints -> universe_constraints
@@ -141,76 +142,70 @@ let to_constraints g s =
 		   "to_constraints: non-trivial algebraic constraint between universes")
   in Constraints.fold tr s Constraint.empty
 
-let eq_constr_univs_infer univs m n =
-  if m == n then true, Constraints.empty
+let eq_constr_univs_infer univs fold m n accu =
+  if m == n then Some accu
   else 
-    let cstrs = ref Constraints.empty in
+    let cstrs = ref accu in
     let eq_universes strict = UGraph.check_eq_instances univs in
     let eq_sorts s1 s2 = 
       if Sorts.equal s1 s2 then true
       else
 	let u1 = Sorts.univ_of_sort s1 and u2 = Sorts.univ_of_sort s2 in
-	  if UGraph.check_eq univs u1 u2 then true
-	  else
-	    (cstrs := Constraints.add (u1, UEq, u2) !cstrs;
-	     true)
+	match fold (Constraints.singleton (u1, UEq, u2)) !cstrs with
+	| None -> false
+	| Some accu -> cstrs := accu; true
     in
     let rec eq_constr' m n = 
       m == n ||	Constr.compare_head_gen eq_universes eq_sorts eq_constr' m n
     in
     let res = Constr.compare_head_gen eq_universes eq_sorts eq_constr' m n in
-      res, !cstrs
+    if res then Some !cstrs else None
 
 (** Variant of [eq_constr_univs_infer] taking kind-of-term functions,
     to expose subterms of [m] and [n], arguments. *)
-let eq_constr_univs_infer_with kind1 kind2 univs m n =
+let eq_constr_univs_infer_with kind1 kind2 univs fold m n accu =
   (* spiwack: duplicates the code of [eq_constr_univs_infer] because I
      haven't find a way to factor the code without destroying
      pointer-equality optimisations in [eq_constr_univs_infer].
      Pointer equality is not sufficient to ensure equality up to
      [kind1,kind2], because [kind1] and [kind2] may be different,
      typically evaluating [m] and [n] in different evar maps. *)
-  let cstrs = ref Constraints.empty in
+  let cstrs = ref accu in
   let eq_universes strict = UGraph.check_eq_instances univs in
   let eq_sorts s1 s2 = 
     if Sorts.equal s1 s2 then true
     else
       let u1 = Sorts.univ_of_sort s1 and u2 = Sorts.univ_of_sort s2 in
-      if UGraph.check_eq univs u1 u2 then true
-      else
-	(cstrs := Constraints.add (u1, UEq, u2) !cstrs;
-	 true)
+      match fold (Constraints.singleton (u1, UEq, u2)) !cstrs with
+      | None -> false
+      | Some accu -> cstrs := accu; true
   in
   let rec eq_constr' m n = 
     Constr.compare_head_gen_with kind1 kind2 eq_universes eq_sorts eq_constr' m n
   in
   let res = Constr.compare_head_gen_with kind1 kind2 eq_universes eq_sorts eq_constr' m n in
-  res, !cstrs
+  if res then Some !cstrs else None
 
-let leq_constr_univs_infer univs m n =
-  if m == n then true, Constraints.empty
+let leq_constr_univs_infer univs fold m n accu =
+  if m == n then Some accu
   else 
-    let cstrs = ref Constraints.empty in
+    let cstrs = ref accu in
     let eq_universes strict l l' = UGraph.check_eq_instances univs l l' in
     let eq_sorts s1 s2 = 
       if Sorts.equal s1 s2 then true
       else
 	let u1 = Sorts.univ_of_sort s1 and u2 = Sorts.univ_of_sort s2 in
-	if UGraph.check_eq univs u1 u2 then true
-	else (cstrs := Constraints.add (u1, UEq, u2) !cstrs;
-	      true)
+	match fold (Constraints.singleton (u1, UEq, u2)) !cstrs with
+	| None -> false
+	| Some accu -> cstrs := accu; true
     in
     let leq_sorts s1 s2 = 
       if Sorts.equal s1 s2 then true
       else 
 	let u1 = Sorts.univ_of_sort s1 and u2 = Sorts.univ_of_sort s2 in
-	  if UGraph.check_leq univs u1 u2 then
-	    ((if Univ.is_small_univ u1 then
-		cstrs := Constraints.add (u1, ULe, u2) !cstrs);
-	     true)
-	  else
-	    (cstrs := Constraints.add (u1, ULe, u2) !cstrs; 
-	     true)
+	match fold (Constraints.singleton (u1, ULe, u2)) !cstrs with
+	| None -> false
+	| Some accu -> cstrs := accu; true
     in
     let rec eq_constr' m n = 
       m == n ||	Constr.compare_head_gen eq_universes eq_sorts eq_constr' m n
@@ -220,7 +215,7 @@ let leq_constr_univs_infer univs m n =
 	eq_constr' leq_constr' m n
     and leq_constr' m n = m == n || compare_leq m n in
     let res = compare_leq m n in
-    res, !cstrs
+    if res then Some !cstrs else None
 
 let eq_constr_universes m n =
   if m == n then true, Constraints.empty
diff --git a/library/universes.mli b/library/universes.mli
index 285580be2..7b17b8898 100644
--- a/library/universes.mli
+++ b/library/universes.mli
@@ -63,6 +63,7 @@ module Constraints : sig
 end
 
 type universe_constraints = Constraints.t
+type 'a constraint_accumulator = universe_constraints -> 'a -> 'a option
 type 'a universe_constrained = 'a * universe_constraints
 type 'a universe_constraint_function = 'a -> 'a -> universe_constraints -> universe_constraints
 
@@ -75,7 +76,8 @@ val to_constraints : UGraph.t -> universe_constraints -> constraints
 
 (** [eq_constr_univs_infer u a b] is [true, c] if [a] equals [b] modulo alpha, casts,
    application grouping, the universe constraints in [u] and additional constraints [c]. *)
-val eq_constr_univs_infer : UGraph.t -> constr -> constr -> bool universe_constrained
+val eq_constr_univs_infer : UGraph.t -> 'a constraint_accumulator ->
+  constr -> constr -> 'a -> 'a option
 
 (** [eq_constr_univs_infer_With kind1 kind2 univs m n] is a variant of
     {!eq_constr_univs_infer} taking kind-of-term functions, to expose
@@ -83,12 +85,13 @@ val eq_constr_univs_infer : UGraph.t -> constr -> constr -> bool universe_constr
 val eq_constr_univs_infer_with :
   (constr -> (constr,types) kind_of_term) ->
   (constr -> (constr,types) kind_of_term) ->
-  UGraph.t -> constr -> constr -> bool universe_constrained
+  UGraph.t -> 'a constraint_accumulator -> constr -> constr -> 'a -> 'a option
 
 (** [leq_constr_univs u a b] is [true, c] if [a] is convertible to [b]
     modulo alpha, casts, application grouping, the universe constraints
     in [u] and additional constraints [c]. *)
-val leq_constr_univs_infer : UGraph.t -> constr -> constr -> bool universe_constrained
+val leq_constr_univs_infer : UGraph.t -> 'a constraint_accumulator ->
+  constr -> constr -> 'a -> 'a option
 
 (** [eq_constr_universes a b] [true, c] if [a] equals [b] modulo alpha, casts,
     application grouping and the universe constraints in [c]. *)
diff --git a/pretyping/evarutil.ml b/pretyping/evarutil.ml
index 508b9e802..3c3afac54 100644
--- a/pretyping/evarutil.ml
+++ b/pretyping/evarutil.ml
@@ -860,13 +860,14 @@ let kind_of_term_upto sigma t =
 let eq_constr_univs_test sigma1 sigma2 t u =
   (* spiwack: mild code duplication with {!Evd.eq_constr_univs}. *)
   let open Evd in
-  let b, c =
+  let fold cstr sigma =
+    try Some (add_universe_constraints sigma cstr)
+    with Univ.UniverseInconsistency _ | UniversesDiffer -> None
+  in
+  let ans =
     Universes.eq_constr_univs_infer_with
       (fun t -> kind_of_term_upto sigma1 t)
       (fun u -> kind_of_term_upto sigma2 u)
-      (universes sigma2) t u
+      (universes sigma2) fold t u sigma2
   in
-  if b then
-    try let _ = add_universe_constraints sigma2 c in true
-    with Univ.UniverseInconsistency _ | UniversesDiffer -> false
-  else false
+  match ans with None -> false | Some _ -> true
diff --git a/pretyping/reductionops.ml b/pretyping/reductionops.ml
index bdd9ed81c..d5a93230f 100644
--- a/pretyping/reductionops.ml
+++ b/pretyping/reductionops.ml
@@ -1299,18 +1299,21 @@ let sigma_univ_state =
 
 let infer_conv_gen conv_fun ?(catch_incon=true) ?(pb=Reduction.CUMUL)
     ?(ts=full_transparent_state) env sigma x y =
-  try 
+  try
+    let fold cstr sigma =
+      try Some (Evd.add_universe_constraints sigma cstr)
+      with Univ.UniverseInconsistency _ | Evd.UniversesDiffer -> None
+    in
     let b, sigma = 
-      let b, cstrs = 
+      let ans =
 	if pb == Reduction.CUMUL then 
-	  Universes.leq_constr_univs_infer (Evd.universes sigma) x y 
+	  Universes.leq_constr_univs_infer (Evd.universes sigma) fold x y sigma
 	else
-	  Universes.eq_constr_univs_infer (Evd.universes sigma) x y 
+	  Universes.eq_constr_univs_infer (Evd.universes sigma) fold x y sigma
       in
-	if b then 
-	  try true, Evd.add_universe_constraints sigma cstrs
-	  with Univ.UniverseInconsistency _ | Evd.UniversesDiffer -> false, sigma
-	else false, sigma
+      match ans with
+      | None -> false, sigma
+      | Some sigma -> true, sigma
     in
       if b then sigma, true
       else