Merge PR #1033: Universe binder improvements

7 files changed, 297 insertions, 120 deletions

diff --git a/engine/evd.ml b/engine/evd.ml
index ecef04204..d57ae89dd 100644
--- a/engine/evd.ml
+++ b/engine/evd.ml
@@ -756,10 +756,12 @@ let evar_universe_context d = d.universes
 
 let universe_context_set d = UState.context_set d.universes
 
-let universe_context ~names ~extensible evd =
-  UState.universe_context ~names ~extensible evd.universes
+let to_universe_context evd = UState.context evd.universes
 
-let check_univ_decl evd decl = UState.check_univ_decl evd.universes decl
+let const_univ_entry ~poly evd = UState.const_univ_entry ~poly evd.universes
+let ind_univ_entry ~poly evd = UState.ind_univ_entry ~poly evd.universes
+
+let check_univ_decl ~poly evd decl = UState.check_univ_decl ~poly evd.universes decl
 
 let restrict_universe_context evd vars =
   { evd with universes = UState.restrict evd.universes vars }
@@ -802,7 +804,7 @@ let make_evar_universe_context e l =
   | Some us ->
       List.fold_left
         (fun uctx (loc,id) ->
-        fst (UState.new_univ_variable ?loc univ_rigid (Some (Id.to_string id)) uctx))
+        fst (UState.new_univ_variable ?loc univ_rigid (Some id) uctx))
         uctx us
 
 (****************************************)
@@ -933,8 +935,7 @@ let nf_constraints evd =
 
 let universe_of_name evd s = UState.universe_of_name evd.universes s
 
-let add_universe_name evd s l =
-  { evd with universes = UState.add_universe_name evd.universes s l }
+let universe_binders evd = UState.universe_binders evd.universes
 
 let universes evd = UState.ugraph evd.universes
 
diff --git a/engine/evd.mli b/engine/evd.mli
index aed2d7083..fb5a6cd16 100644
--- a/engine/evd.mli
+++ b/engine/evd.mli
@@ -493,6 +493,8 @@ type 'a in_evar_universe_context = 'a * UState.t
 val evar_universe_context_set : UState.t -> Univ.ContextSet.t
 val evar_universe_context_constraints : UState.t -> Univ.constraints
 val evar_context_universe_context : UState.t -> Univ.UContext.t
+[@@ocaml.deprecated "alias of UState.context"]
+
 val evar_universe_context_of : Univ.ContextSet.t -> UState.t
 val empty_evar_universe_context : UState.t
 val union_evar_universe_context : UState.t -> UState.t ->
@@ -503,14 +505,14 @@ val constrain_variables : Univ.LSet.t -> UState.t -> UState.t
 
 val evar_universe_context_of_binders :
   Universes.universe_binders -> UState.t
-							    
+
 val make_evar_universe_context : env -> (Id.t located) list option -> UState.t
-val restrict_universe_context : evar_map -> Univ.LSet.t -> evar_map							   
+val restrict_universe_context : evar_map -> Univ.LSet.t -> evar_map
 (** Raises Not_found if not a name for a universe in this map. *)
-val universe_of_name : evar_map -> string -> Univ.Level.t
-val add_universe_name : evar_map -> string -> Univ.Level.t -> evar_map
+val universe_of_name : evar_map -> Id.t -> Univ.Level.t
 
-val add_constraints_context : UState.t -> 
+val universe_binders : evar_map -> Universes.universe_binders
+val add_constraints_context : UState.t ->
   Univ.constraints -> UState.t
 
 
@@ -520,9 +522,9 @@ val normalize_evar_universe_context_variables : UState.t ->
 val normalize_evar_universe_context : UState.t -> 
   UState.t
 
-val new_univ_level_variable : ?loc:Loc.t -> ?name:string -> rigid -> evar_map -> evar_map * Univ.Level.t
-val new_univ_variable : ?loc:Loc.t -> ?name:string -> rigid -> evar_map -> evar_map * Univ.Universe.t
-val new_sort_variable : ?loc:Loc.t -> ?name:string -> rigid -> evar_map -> evar_map * Sorts.t
+val new_univ_level_variable : ?loc:Loc.t -> ?name:Id.t -> rigid -> evar_map -> evar_map * Univ.Level.t
+val new_univ_variable : ?loc:Loc.t -> ?name:Id.t -> rigid -> evar_map -> evar_map * Univ.Universe.t
+val new_sort_variable : ?loc:Loc.t -> ?name:Id.t -> rigid -> evar_map -> evar_map * Sorts.t
 
 val add_global_univ : evar_map -> Univ.Level.t -> evar_map
 
@@ -552,13 +554,20 @@ val check_leq : evar_map -> Univ.Universe.t -> Univ.Universe.t -> bool
 
 val evar_universe_context : evar_map -> UState.t
 val universe_context_set : evar_map -> Univ.ContextSet.t
-val universe_context : names:(Id.t located) list -> extensible:bool -> evar_map ->
-		       (Id.t * Univ.Level.t) list * Univ.UContext.t
 val universe_subst : evar_map -> Universes.universe_opt_subst
 val universes : evar_map -> UGraph.t
 
-val check_univ_decl : evar_map -> UState.universe_decl ->
-  Universes.universe_binders * Univ.UContext.t
+(** [to_universe_context evm] extracts the local universes and
+    constraints of [evm] and orders the universes the same as
+    [Univ.ContextSet.to_context]. *)
+val to_universe_context : evar_map -> Univ.UContext.t
+
+val const_univ_entry : poly:bool -> evar_map -> Entries.constant_universes_entry
+
+(** NB: [ind_univ_entry] cannot create cumulative entries. *)
+val ind_univ_entry : poly:bool -> evar_map -> Entries.inductive_universes
+
+val check_univ_decl : poly:bool -> evar_map -> UState.universe_decl -> Entries.constant_universes_entry
 
 val merge_universe_context : evar_map -> UState.t -> evar_map
 val set_universe_context : evar_map -> UState.t -> evar_map
diff --git a/engine/uState.ml b/engine/uState.ml
index 01a479821..4e30640e4 100644
--- a/engine/uState.ml
+++ b/engine/uState.ml
@@ -11,28 +11,16 @@ open CErrors
 open Util
 open Names
 
-module StringOrd = struct type t = string let compare = String.compare end
-module UNameMap = struct
+module UNameMap = Names.Id.Map
     
-  include Map.Make(StringOrd)
-    
-  let union s t = 
-    if s == t then s
-    else
-      merge (fun k l r -> 
-        match l, r with
-        | Some _, _ -> l
-        | _, _ -> r) s t
-end
-
 type uinfo = {
-  uname : string option;
+  uname : Id.t option;
   uloc : Loc.t option;
 }
 
 (* 2nd part used to check consistency on the fly. *)
 type t =
- { uctx_names : Univ.Level.t UNameMap.t * uinfo Univ.LMap.t;
+ { uctx_names : Universes.universe_binders * uinfo Univ.LMap.t;
    uctx_local : Univ.ContextSet.t; (** The local context of variables *)
    uctx_univ_variables : Universes.universe_opt_subst;
    (** The local universes that are unification variables *)
@@ -59,12 +47,20 @@ let is_empty ctx =
   Univ.ContextSet.is_empty ctx.uctx_local && 
     Univ.LMap.is_empty ctx.uctx_univ_variables
 
+let uname_union s t =
+  if s == t then s
+  else
+    UNameMap.merge (fun k l r ->
+        match l, r with
+        | Some _, _ -> l
+        | _, _ -> r) s t
+
 let union ctx ctx' =
   if ctx == ctx' then ctx
   else if is_empty ctx' then ctx
   else
     let local = Univ.ContextSet.union ctx.uctx_local ctx'.uctx_local in
-    let names = UNameMap.union (fst ctx.uctx_names) (fst ctx'.uctx_names) in
+    let names = uname_union (fst ctx.uctx_names) (fst ctx'.uctx_names) in
     let newus = Univ.LSet.diff (Univ.ContextSet.levels ctx'.uctx_local)
                                (Univ.ContextSet.levels ctx.uctx_local) in
     let newus = Univ.LSet.diff newus (Univ.LMap.domain ctx.uctx_univ_variables) in
@@ -91,6 +87,17 @@ let constraints ctx = snd ctx.uctx_local
 
 let context ctx = Univ.ContextSet.to_context ctx.uctx_local
 
+let const_univ_entry ~poly uctx =
+  let open Entries in
+  if poly then Polymorphic_const_entry (context uctx)
+  else Monomorphic_const_entry (context_set uctx)
+
+(* does not support cumulativity since you need more info *)
+let ind_univ_entry ~poly uctx =
+  let open Entries in
+  if poly then Polymorphic_ind_entry (context uctx)
+  else Monomorphic_ind_entry (context_set uctx)
+
 let of_context_set ctx = { empty with uctx_local = ctx }
 
 let subst ctx = ctx.uctx_univ_variables
@@ -102,6 +109,9 @@ let initial_graph ctx = ctx.uctx_initial_universes
 let algebraics ctx = ctx.uctx_univ_algebraic
 
 let add_uctx_names ?loc s l (names, names_rev) =
+  if UNameMap.mem s names
+  then user_err ?loc ~hdr:"add_uctx_names"
+      Pp.(str "Universe " ++ Names.Id.print s ++ str" already bound.");
   (UNameMap.add s l names, Univ.LMap.add l { uname = Some s; uloc = loc } names_rev)
 
 let add_uctx_loc l loc (names, names_rev) =
@@ -111,10 +121,14 @@ let add_uctx_loc l loc (names, names_rev) =
 
 let of_binders b =
   let ctx = empty in
-  let names =
-    List.fold_left (fun acc (id, l) -> add_uctx_names (Id.to_string id) l acc)
-                   ctx.uctx_names b
-  in { ctx with uctx_names = names }
+  let rmap =
+    UNameMap.fold (fun id l rmap ->
+        Univ.LMap.add l { uname = Some id; uloc = None } rmap)
+      b Univ.LMap.empty
+  in
+  { ctx with uctx_names = b, rmap }
+
+let universe_binders ctx = fst ctx.uctx_names
 
 let instantiate_variable l b v =
   try v := Univ.LMap.update l (Some b) !v
@@ -253,69 +267,105 @@ let constrain_variables diff ctx =
 let pr_uctx_level uctx = 
   let map, map_rev = uctx.uctx_names in 
     fun l ->
-      try str (Option.get (Univ.LMap.find l map_rev).uname)
+      try Id.print (Option.get (Univ.LMap.find l map_rev).uname)
       with Not_found | Option.IsNone ->
         Universes.pr_with_global_universes l
 
 type universe_decl =
   (Names.Id.t Loc.located list, Univ.Constraint.t) Misctypes.gen_universe_decl
 
-let universe_context ~names ~extensible ctx =
-  let levels = Univ.ContextSet.levels ctx.uctx_local in
+let error_unbound_universes left uctx =
+  let open Univ in
+  let n = LSet.cardinal left in
+  let loc =
+    try
+      let info =
+        LMap.find (LSet.choose left) (snd uctx.uctx_names) in
+      info.uloc
+    with Not_found -> None
+  in
+  user_err ?loc ~hdr:"universe_context"
+    ((str(CString.plural n "Universe") ++ spc () ++
+      LSet.pr (pr_uctx_level uctx) left ++
+      spc () ++ str (CString.conjugate_verb_to_be n) ++
+      str" unbound."))
+
+let universe_context ~names ~extensible uctx =
+  let open Univ in
+  let levels = ContextSet.levels uctx.uctx_local in
   let newinst, left =
     List.fold_right
       (fun (loc,id) (newinst, acc) ->
          let l =
-           try UNameMap.find (Id.to_string id) (fst ctx.uctx_names)
-           with Not_found ->
-             user_err ?loc ~hdr:"universe_context"
-               (str"Universe " ++ Id.print id ++ str" is not bound anymore.")
-         in (l :: newinst, Univ.LSet.remove l acc))
+           try UNameMap.find id (fst uctx.uctx_names)
+           with Not_found -> assert false
+         in (l :: newinst, LSet.remove l acc))
       names ([], levels)
   in
-  if not extensible && not (Univ.LSet.is_empty left) then
-    let n = Univ.LSet.cardinal left in
-    let loc =
-      try
-        let info =
-          Univ.LMap.find (Univ.LSet.choose left) (snd ctx.uctx_names) in
-        info.uloc
-      with Not_found -> None
-    in
-    user_err ?loc ~hdr:"universe_context"
-      ((str(CString.plural n "Universe") ++ spc () ++
-	Univ.LSet.pr (pr_uctx_level ctx) left ++
-	spc () ++ str (CString.conjugate_verb_to_be n) ++
-        str" unbound."))
+  if not extensible && not (LSet.is_empty left)
+  then error_unbound_universes left uctx
   else
-    let left = Univ.ContextSet.sort_levels (Array.of_list (Univ.LSet.elements left)) in
+    let left = ContextSet.sort_levels (Array.of_list (LSet.elements left)) in
     let inst = Array.append (Array.of_list newinst) left in
-    let inst = Univ.Instance.of_array inst in
-    let map = List.map (fun (s,l) -> Id.of_string s, l) (UNameMap.bindings (fst ctx.uctx_names)) in
-    let ctx = Univ.UContext.make (inst,
-                                  Univ.ContextSet.constraints ctx.uctx_local) in
-    map, ctx
+    let inst = Instance.of_array inst in
+    let ctx = UContext.make (inst, ContextSet.constraints uctx.uctx_local) in
+    ctx
+
+let check_universe_context_set ~names ~extensible uctx =
+  if extensible then ()
+  else
+    let open Univ in
+    let left = List.fold_left (fun left (loc,id) ->
+        let l =
+          try UNameMap.find id (fst uctx.uctx_names)
+          with Not_found -> assert false
+        in LSet.remove l left)
+        (ContextSet.levels uctx.uctx_local) names
+    in
+    if not (LSet.is_empty left)
+    then error_unbound_universes left uctx
 
-let check_implication uctx cstrs ctx =
+let check_implication uctx cstrs cstrs' =
   let gr = initial_graph uctx in
   let grext = UGraph.merge_constraints cstrs gr in
-  let cstrs' = Univ.UContext.constraints ctx in
   if UGraph.check_constraints cstrs' grext then ()
   else CErrors.user_err ~hdr:"check_univ_decl"
       (str "Universe constraints are not implied by the ones declared.")
 
-let check_univ_decl uctx decl =
+let check_mono_univ_decl uctx decl =
   let open Misctypes in
-  let pl, ctx = universe_context
-      ~names:decl.univdecl_instance
-      ~extensible:decl.univdecl_extensible_instance
-      uctx
+  let () =
+    let names = decl.univdecl_instance in
+    let extensible = decl.univdecl_extensible_instance in
+    check_universe_context_set ~names ~extensible uctx
   in
   if not decl.univdecl_extensible_constraints then
-    check_implication uctx decl.univdecl_constraints ctx;
-  pl, ctx
+    check_implication uctx
+      decl.univdecl_constraints
+      (Univ.ContextSet.constraints uctx.uctx_local);
+  uctx.uctx_local
+
+let check_univ_decl ~poly uctx decl =
+  let open Misctypes in
+  let ctx =
+    let names = decl.univdecl_instance in
+    let extensible = decl.univdecl_extensible_instance in
+    if poly
+    then Entries.Polymorphic_const_entry (universe_context ~names ~extensible uctx)
+    else
+      let () = check_universe_context_set ~names ~extensible uctx in
+      Entries.Monomorphic_const_entry uctx.uctx_local
+  in
+  if not decl.univdecl_extensible_constraints then
+    check_implication uctx
+      decl.univdecl_constraints
+      (Univ.ContextSet.constraints uctx.uctx_local);
+  ctx
 
 let restrict ctx vars =
+  let vars = Names.Id.Map.fold (fun na l vars -> Univ.LSet.add l vars)
+      (fst ctx.uctx_names) vars
+  in
   let uctx' = Univops.restrict_universe_context ctx.uctx_local vars in
   { ctx with uctx_local = uctx' }
 
@@ -526,10 +576,6 @@ let normalize uctx =
 let universe_of_name uctx s = 
   UNameMap.find s (fst uctx.uctx_names)
 
-let add_universe_name uctx s l =
-  let names' = add_uctx_names s l uctx.uctx_names in
-  { uctx with uctx_names = names' }
-
 let update_sigma_env uctx env =
   let univs = Environ.universes env in
   let eunivs =
diff --git a/engine/uState.mli b/engine/uState.mli
index 1c906fcb2..16fba41e0 100644
--- a/engine/uState.mli
+++ b/engine/uState.mli
@@ -32,6 +32,8 @@ val of_context_set : Univ.ContextSet.t -> t
 
 val of_binders : Universes.universe_binders -> t
 
+val universe_binders : t -> Universes.universe_binders
+
 (** {5 Projections} *)
 
 val context_set : t -> Univ.ContextSet.t
@@ -57,6 +59,13 @@ val constraints : t -> Univ.constraints
 val context : t -> Univ.UContext.t
 (** Shorthand for {!context_set} with {!Context_set.to_context}. *)
 
+val const_univ_entry : poly:bool -> t -> Entries.constant_universes_entry
+(** Pick from {!context} or {!context_set} based on [poly]. *)
+
+val ind_univ_entry : poly:bool -> t -> Entries.inductive_universes
+(** Pick from {!context} or {!context_set} based on [poly].
+    Cannot create cumulative entries. *)
+
 (** {5 Constraints handling} *)
 
 val add_constraints : t -> Univ.constraints -> t
@@ -71,10 +80,7 @@ val add_universe_constraints : t -> Universes.universe_constraints -> t
 
 (** {5 Names} *)
 
-val add_universe_name : t -> string -> Univ.Level.t -> t
-(** Associate a human-readable name to a local variable. *)
-
-val universe_of_name : t -> string -> Univ.Level.t
+val universe_of_name : t -> Id.t -> Univ.Level.t
 (** Retrieve the universe associated to the name. *)
 
 (** {5 Unification} *)
@@ -93,7 +99,7 @@ val merge : ?loc:Loc.t -> bool -> rigid -> t -> Univ.ContextSet.t -> t
 val merge_subst : t -> Universes.universe_opt_subst -> t
 val emit_side_effects : Safe_typing.private_constants -> t -> t
 
-val new_univ_variable : ?loc:Loc.t -> rigid -> string option -> t -> t * Univ.Level.t
+val new_univ_variable : ?loc:Loc.t -> rigid -> Id.t option -> t -> t * Univ.Level.t
 val add_global_univ : t -> Univ.Level.t -> t
 
 (** [make_flexible_variable g algebraic l]
@@ -123,24 +129,23 @@ val refresh_undefined_univ_variables : t -> t * Univ.universe_level_subst
 
 val normalize : t -> t
 
-(** [universe_context names extensible ctx]
+type universe_decl =
+  (Names.Id.t Loc.located list, Univ.Constraint.t) Misctypes.gen_universe_decl
 
-    Return a universe context containing the local universes of [ctx]
-    and their constraints. The universes corresponding to [names] come
-    first in the order defined by that list.
+(** [check_univ_decl ctx decl]
 
-    If [extensible] is false, check that the universes of [names] are
-    the only local universes.
+   If non extensible in [decl], check that the local universes (resp.
+   universe constraints) in [ctx] are implied by [decl].
 
-    Also return the association list of universe names and universes
-    (including those not in [names]). *)
-val universe_context : names:(Id.t Loc.located) list -> extensible:bool -> t ->
-  (Id.t * Univ.Level.t) list * Univ.UContext.t
+   Return a [Entries.constant_universes_entry] containing the local
+   universes of [ctx] and their constraints.
 
-type universe_decl =
-  (Names.Id.t Loc.located list, Univ.Constraint.t) Misctypes.gen_universe_decl
+   When polymorphic, the universes corresponding to
+   [decl.univdecl_instance] come first in the order defined by that
+   list. *)
+val check_univ_decl : poly:bool -> t -> universe_decl -> Entries.constant_universes_entry
 
-val check_univ_decl : t -> universe_decl -> Universes.universe_binders * Univ.UContext.t
+val check_mono_univ_decl : t -> universe_decl -> Univ.ContextSet.t
 
 (** {5 TODO: Document me} *)
 
diff --git a/engine/universes.ml b/engine/universes.ml
index 6c1b64d74..5ac1bc685 100644
--- a/engine/universes.ml
+++ b/engine/universes.ml
@@ -21,18 +21,63 @@ let pr_with_global_universes l =
 
 (** Local universe names of polymorphic references *)
 
-type universe_binders = (Id.t * Univ.Level.t) list
+type universe_binders = Univ.Level.t Names.Id.Map.t
+
+let empty_binders = Id.Map.empty
 
 let universe_binders_table = Summary.ref Refmap.empty ~name:"universe binders"
 
-let universe_binders_of_global ref =
+let universe_binders_of_global ref : universe_binders =
   try
     let l = Refmap.find ref !universe_binders_table in l
-  with Not_found -> []
+  with Not_found -> Names.Id.Map.empty
 
-let register_universe_binders ref l =
+let cache_ubinder (_,(ref,l)) =
   universe_binders_table := Refmap.add ref l !universe_binders_table
 
+let subst_ubinder (subst,(ref,l as orig)) =
+  let ref' = fst (Globnames.subst_global subst ref) in
+  if ref == ref' then orig else ref', l
+
+let discharge_ubinder (_,(ref,l)) =
+  Some (Lib.discharge_global ref, l)
+
+let ubinder_obj : Globnames.global_reference * universe_binders -> Libobject.obj =
+  let open Libobject in
+  declare_object { (default_object "universe binder") with
+    cache_function = cache_ubinder;
+    load_function = (fun _ x -> cache_ubinder x);
+    classify_function = (fun x -> Substitute x);
+    subst_function = subst_ubinder;
+    discharge_function = discharge_ubinder;
+    rebuild_function = (fun x -> x); }
+
+let register_universe_binders ref ubinders =
+  (* Add the polymorphic (section) universes *)
+  let open Names in
+  let ubinders = Id.Map.fold (fun id (poly,lvl) ubinders ->
+      if poly then Id.Map.add id lvl ubinders
+      else ubinders)
+      (fst (Global.global_universe_names ())) ubinders
+  in
+  if not (Id.Map.is_empty ubinders)
+  then Lib.add_anonymous_leaf (ubinder_obj (ref,ubinders))
+
+type univ_name_list = Name.t Loc.located list
+
+let universe_binders_with_opt_names ref levels = function
+  | None -> universe_binders_of_global ref
+  | Some udecl ->
+    if Int.equal(List.length levels) (List.length udecl)
+    then
+      List.fold_left2 (fun acc (_,na) lvl -> match na with
+          | Anonymous -> acc
+          | Name na -> Names.Id.Map.add na lvl acc)
+        empty_binders udecl levels
+    else
+      CErrors.user_err ~hdr:"universe_binders_with_opt_names"
+        Pp.(str "Universe instance should have length " ++ int (List.length levels))
+
 (* To disallow minimization to Set *)
 
 let set_minimization = ref true
diff --git a/engine/universes.mli b/engine/universes.mli
index a960099ed..1401c4ee8 100644
--- a/engine/universes.mli
+++ b/engine/universes.mli
@@ -21,11 +21,24 @@ val pr_with_global_universes : Level.t -> Pp.t
 
 (** Local universe name <-> level mapping *)
 
-type universe_binders = (Id.t * Univ.Level.t) list
+type universe_binders = Univ.Level.t Names.Id.Map.t
+
+val empty_binders : universe_binders
 
 val register_universe_binders : Globnames.global_reference -> universe_binders -> unit
 val universe_binders_of_global : Globnames.global_reference -> universe_binders
 
+type univ_name_list = Name.t Loc.located list
+
+(** [universe_binders_with_opt_names ref u l]
+
+    If [l] is [Some univs] return the universe binders naming the levels of [u] by [univs] (skipping Anonymous).
+    May error if the lengths mismatch.
+
+    Otherwise return [universe_binders_of_global ref]. *)
+val universe_binders_with_opt_names : Globnames.global_reference ->
+  Univ.Level.t list -> univ_name_list option -> universe_binders
+
 (** The global universe counter *)
 val set_remote_new_univ_level : Level.t RemoteCounter.installer
 
diff --git a/engine/univops.ml b/engine/univops.ml
index d498b2e0d..9a9ae12ca 100644
--- a/engine/univops.ml
+++ b/engine/univops.ml
@@ -20,21 +20,79 @@ let universes_of_constr c =
     | _ -> Constr.fold aux s c
   in aux LSet.empty c
 
-let restrict_universe_context (univs,csts) s =
-  (* Universes that are not necessary to typecheck the term.
-     E.g. univs introduced by tactics and not used in the proof term. *)
-  let diff = LSet.diff univs s in
-  let rec aux diff candid univs ness = 
-    let (diff', candid', univs', ness') = 
-      Constraint.fold
-	(fun (l, d, r as c) (diff, candid, univs, csts) ->
-	  if not (LSet.mem l diff) then
-	    (LSet.remove r diff, candid, univs, Constraint.add c csts)
-	  else if not (LSet.mem r diff) then
-	    (LSet.remove l diff, candid, univs, Constraint.add c csts)
-	  else (diff, Constraint.add c candid, univs, csts))
-	candid (diff, Constraint.empty, univs, ness)
-    in
-      if ness' == ness then (LSet.diff univs diff', ness)
-      else aux diff' candid' univs' ness'
-  in aux diff csts univs Constraint.empty
+type graphnode = {
+  mutable up : constraint_type LMap.t;
+  mutable visited : bool
+}
+
+let merge_types d d0 =
+  match d, d0 with
+  | _, Lt | Lt, _ -> Lt
+  | Le, _ | _, Le -> Le
+  | Eq, Eq -> Eq
+
+let merge_up d b up =
+  let find = try Some (LMap.find b up) with Not_found -> None in
+  match find with
+  | Some d0 ->
+    let d = merge_types d d0 in
+    if d == d0 then up else LMap.add b d up
+  | None -> LMap.add b d up
+
+let add_up a d b graph =
+  let node, graph =
+    try LMap.find a graph, graph
+    with Not_found ->
+      let node = { up = LMap.empty; visited = false } in
+      node, LMap.add a node graph
+  in
+  node.up <- merge_up d b node.up;
+  graph
+
+(* for each node transitive close until you find a non removable, discard the rest *)
+let transitive_close removable graph =
+  let rec do_node a node =
+    if not node.visited
+    then
+      let keepup =
+        LMap.fold (fun b d keepup ->
+            if not (LSet.mem b removable)
+            then merge_up d b keepup
+            else
+              begin
+                match LMap.find b graph with
+                | bnode ->
+                  do_node b bnode;
+                  LMap.fold (fun k d' keepup ->
+                    merge_up (merge_types d d') k keepup)
+                    bnode.up keepup
+                | exception Not_found -> keepup
+              end
+          )
+          node.up LMap.empty
+      in
+      node.up <- keepup;
+      node.visited <- true
+  in
+  LMap.iter do_node graph
+
+let restrict_universe_context (univs,csts) keep =
+  let removable = LSet.diff univs keep in
+  let (csts, rem) =
+    Constraint.fold (fun (a,d,b as cst) (csts, rem) ->
+        if LSet.mem a removable || LSet.mem b removable
+        then (csts, add_up a d b rem)
+        else (Constraint.add cst csts, rem))
+      csts (Constraint.empty, LMap.empty)
+  in
+  transitive_close removable rem;
+  let csts =
+    LMap.fold (fun a node csts ->
+        if LSet.mem a removable
+        then csts
+        else
+          LMap.fold (fun b d csts -> Constraint.add (a,d,b) csts)
+            node.up csts)
+      rem csts
+  in
+  (LSet.inter univs keep, csts)