1 files changed, 32 insertions, 22 deletions

diff --git a/engine/uState.ml b/engine/uState.ml
index 5747ae2ad4..6aecc368e6 100644
--- a/engine/uState.ml
+++ b/engine/uState.ml
@@ -441,11 +441,13 @@ let restrict_universe_context (univs, csts) keep =
   if LSet.is_empty removed then univs, csts
   else
   let allunivs = Constraint.fold (fun (u,_,v) all -> LSet.add u (LSet.add v all)) csts univs in
-  let g = UGraph.empty_universes in
-  let g = LSet.fold UGraph.add_universe_unconstrained allunivs g in
+  let g = UGraph.initial_universes in
+  let g = LSet.fold (fun v g -> if Level.is_small v then g else UGraph.add_universe v false g) allunivs g in
   let g = UGraph.merge_constraints csts g in
-  let allkept = LSet.diff allunivs removed in
+  let allkept = LSet.union (UGraph.domain UGraph.initial_universes) (LSet.diff allunivs removed) in
   let csts = UGraph.constraints_for ~kept:allkept g in
+  let csts = Constraint.filter (fun (l,d,r) ->
+      not ((Level.is_set l && d == Le) || (Level.is_prop l && d == Lt && Level.is_set r))) csts in
   (LSet.inter univs keep, csts)
 
 let restrict ctx vars =
@@ -575,25 +577,33 @@ let add_global_univ uctx u =
                                      uctx_universes = univs }
 
 let make_flexible_variable ctx ~algebraic u =
-  let {uctx_local = cstrs; uctx_univ_variables = uvars; uctx_univ_algebraic = avars} = ctx in
-  let uvars' = Univ.LMap.add u None uvars in
-  let avars' = 
-    if algebraic then
-      let uu = Univ.Universe.make u in
-      let substu_not_alg u' v =
-        Option.cata (fun vu -> Univ.Universe.equal uu vu && not (Univ.LSet.mem u' avars)) false v
-      in
-      let has_upper_constraint () =
-        Univ.Constraint.exists
-          (fun (l,d,r) -> d == Univ.Lt && Univ.Level.equal l u)
-          (Univ.ContextSet.constraints cstrs)
-      in
-        if not (Univ.LMap.exists substu_not_alg uvars || has_upper_constraint ())
-        then Univ.LSet.add u avars else avars 
-    else avars 
-  in
-  {ctx with uctx_univ_variables = uvars'; 
-      uctx_univ_algebraic = avars'}
+  let open Univ in
+  let {uctx_local = cstrs; uctx_univ_variables = uvars;
+       uctx_univ_algebraic = avars; uctx_universes=g; } = ctx in
+  assert (try LMap.find u uvars == None with Not_found -> true);
+  match UGraph.choose (fun v -> not (Level.equal u v) && (algebraic || not (LSet.mem v avars))) g u with
+  | Some v ->
+    let uvars' = LMap.add u (Some (Universe.make v)) uvars in
+    { ctx with uctx_univ_variables = uvars'; }
+  | None ->
+    let uvars' = LMap.add u None uvars in
+    let avars' =
+      if algebraic then
+        let uu = Universe.make u in
+        let substu_not_alg u' v =
+          Option.cata (fun vu -> Universe.equal uu vu && not (LSet.mem u' avars)) false v
+        in
+        let has_upper_constraint () =
+          Constraint.exists
+            (fun (l,d,r) -> d == Lt && Level.equal l u)
+            (ContextSet.constraints cstrs)
+        in
+        if not (LMap.exists substu_not_alg uvars || has_upper_constraint ())
+        then LSet.add u avars else avars
+      else avars
+    in
+    {ctx with uctx_univ_variables = uvars';
+              uctx_univ_algebraic = avars'}
 
 let make_nonalgebraic_variable ctx u =
   { ctx with uctx_univ_algebraic = Univ.LSet.remove u ctx.uctx_univ_algebraic }