Fix #5347: unify declaration of axioms with and without bound univs.

Note that this makes the following syntax valid:

     Axiom foo@{i} bar : Type@{i}.

(ie putting a universe declaration on the first axiom in the list, the
declaration then holds for the whole list).

3 files changed, 77 insertions, 80 deletions

diff --git a/vernac/classes.ml b/vernac/classes.ml
index 1833b7a1d..b80741269 100644
--- a/vernac/classes.ml
+++ b/vernac/classes.ml
@@ -412,15 +412,15 @@ let context poly l =
       in
       let impl = List.exists test impls in
       let decl = (Discharge, poly, Definitional) in
+      let univs = if poly
+        then Polymorphic_const_entry (Univ.ContextSet.to_context !uctx)
+        else Monomorphic_const_entry !uctx
+      in
       let nstatus = match b with
       | None ->
-        pi3 (Command.declare_assumption false decl (t, !uctx) Universes.empty_binders [] impl
+        pi3 (Command.declare_assumption false decl (t, univs) Universes.empty_binders [] impl
           Vernacexpr.NoInline (Loc.tag id))
       | Some b ->
-        let univs = if poly
-          then Polymorphic_const_entry (Univ.ContextSet.to_context !uctx)
-          else Monomorphic_const_entry !uctx
-        in
         let decl = (Discharge, poly, Definition) in
         let entry = Declare.definition_entry ~univs ~types:t b in
         let hook = Lemmas.mk_hook (fun _ gr -> gr) in
diff --git a/vernac/command.ml b/vernac/command.ml
index 1a8b16a89..373e5a1be 100644
--- a/vernac/command.ml
+++ b/vernac/command.ml
@@ -177,6 +177,10 @@ let do_definition ident k univdecl bl red_option c ctypopt hook =
 let declare_assumption is_coe (local,p,kind) (c,ctx) pl imps impl nl (_,ident) =
 match local with
 | Discharge when Lib.sections_are_opened () ->
+  let ctx = match ctx with
+    | Monomorphic_const_entry ctx -> ctx
+    | Polymorphic_const_entry ctx -> Univ.ContextSet.of_context ctx
+  in
   let decl = (Lib.cwd(), SectionLocalAssum ((c,ctx),p,impl), IsAssumption kind) in
   let _ = declare_variable ident decl in
   let () = assumption_message ident in
@@ -197,10 +201,6 @@ match local with
     | DefaultInline -> Some (Flags.get_inline_level())
     | InlineAt i -> Some i
   in
-  let ctx = if p
-    then Polymorphic_const_entry (Univ.ContextSet.to_context ctx)
-    else Monomorphic_const_entry ctx
-  in
   let decl = (ParameterEntry (None,(c,ctx),inl), IsAssumption kind) in
   let kn = declare_constant ident ~local decl in
   let gr = ConstRef kn in
@@ -221,26 +221,63 @@ let interp_assumption evdref env impls bl c =
   let ty = EConstr.Unsafe.to_constr ty in
   (ty, impls)
 
-let declare_assumptions idl is_coe k (c,ctx) pl imps impl_is_on nl =
+(* When monomorphic the universe constraints are declared with the first declaration only. *)
+let next_uctx =
+  let empty_uctx = Monomorphic_const_entry Univ.ContextSet.empty in
+  function
+  | Polymorphic_const_entry _ as uctx -> uctx
+  | Monomorphic_const_entry _ -> empty_uctx
+
+let declare_assumptions idl is_coe k (c,uctx) pl imps nl =
   let refs, status, _ =
-    List.fold_left (fun (refs,status,ctx) id ->
+    List.fold_left (fun (refs,status,uctx) id ->
       let ref',u',status' =
-	declare_assumption is_coe k (c,ctx) pl imps impl_is_on nl id in
-      (ref',u')::refs, status' && status, Univ.ContextSet.empty)
-      ([],true,ctx) idl
+        declare_assumption is_coe k (c,uctx) pl imps false nl id in
+      (ref',u')::refs, status' && status, next_uctx uctx)
+      ([],true,uctx) idl
   in
   List.rev refs, status
 
-let do_assumptions_unbound_univs (_, poly, _ as kind) nl l =
+
+let maybe_error_many_udecls = function
+  | ((loc,id), Some _) ->
+    user_err ?loc ~hdr:"many_universe_declarations"
+      Pp.(str "When declaring multiple axioms in one command, " ++
+          str "only the first is allowed a universe binder " ++
+          str "(which will be shared by the whole block).")
+  | (_, None) -> ()
+
+let process_assumptions_udecls kind l =
+  let udecl, first_id = match l with
+    | (coe, ((id, udecl)::rest, c))::rest' ->
+      List.iter maybe_error_many_udecls rest;
+      List.iter (fun (coe, (idl, c)) -> List.iter maybe_error_many_udecls idl) rest';
+      udecl, id
+    | (_, ([], _))::_ | [] -> assert false
+  in
+  let () = match kind, udecl with
+    | (Discharge, _, _), Some _ when Lib.sections_are_opened () ->
+      let loc = fst first_id in
+      let msg = Pp.str "Section variables cannot be polymorphic." in
+      user_err ?loc  msg
+    | _ -> ()
+  in
+  udecl, List.map (fun (coe, (idl, c)) -> coe, (List.map fst idl, c)) l
+
+let do_assumptions kind nl l =
   let open Context.Named.Declaration in
   let env = Global.env () in
-  let evdref = ref (Evd.from_env env) in
-  let l = 
-    if poly then
+  let udecl, l = process_assumptions_udecls kind l in
+  let evdref, udecl =
+    let evd, udecl = Univdecls.interp_univ_decl_opt env udecl in
+    ref evd, udecl
+  in
+  let l =
+    if pi2 kind (* poly *) then
       (* Separate declarations so that A B : Type puts A and B in different levels. *)
       List.fold_right (fun (is_coe,(idl,c)) acc ->
-        List.fold_right (fun id acc -> 
-	  (is_coe, ([id], c)) :: acc) idl acc)
+        List.fold_right (fun id acc ->
+          (is_coe, ([id], c)) :: acc) idl acc)
         l []
     else l
   in
@@ -252,71 +289,31 @@ let do_assumptions_unbound_univs (_, poly, _ as kind) nl l =
     let ienv = List.fold_right (fun (_,id) ienv ->
       let impls = compute_internalization_data env Variable t imps in
       Id.Map.add id impls ienv) idl ienv in
-      ((env,ienv),((is_coe,idl),t,imps))) 
+      ((env,ienv),((is_coe,idl),t,imps)))
     (env,empty_internalization_env) l
   in
   let evd = solve_remaining_evars all_and_fail_flags env !evdref Evd.empty in
   (* The universe constraints come from the whole telescope. *)
   let evd = Evd.nf_constraints evd in
-  let ctx = Evd.universe_context_set evd in
-  let nf_evar c = EConstr.Unsafe.to_constr (nf_evar evd (EConstr.of_constr c)) in
-  let l = List.map (on_pi2 nf_evar) l in
-  pi2 (List.fold_left (fun (subst,status,ctx) ((is_coe,idl),t,imps) ->
-    let t = replace_vars subst t in
-    let (refs,status') = declare_assumptions idl is_coe kind (t,ctx) Universes.empty_binders imps false nl in
-    let subst' = List.map2 
-      (fun (_,id) (c,u) -> (id,Universes.constr_of_global_univ (c,u)))
-      idl refs 
-    in
-    (subst'@subst, status' && status,
-     (* The universe constraints are declared with the first declaration only. *)
-     Univ.ContextSet.empty)) ([],true,ctx) l)
-
-let do_assumptions_bound_univs coe kind nl id pl c =
-  let env = Global.env () in
-  let evd, decl = Univdecls.interp_univ_decl_opt env pl in
-  let evdref = ref evd in
-  let ty, impls = interp_type_evars_impls env evdref c in
-  let nf, subst = Evarutil.e_nf_evars_and_universes evdref in
-  let ty = EConstr.Unsafe.to_constr ty in
-  let ty = nf ty in
-  let vars = Univops.universes_of_constr ty in
-  let evd = Evd.restrict_universe_context !evdref vars in
-  let uctx = Evd.check_univ_decl ~poly:(pi2 kind) evd decl in
-  let binders = Evd.universe_binders evd in
-  let uctx = match uctx with
-    | Polymorphic_const_entry uctx ->
-      (* ?? *)
-      Univ.ContextSet.of_context uctx
-    | Monomorphic_const_entry uctx -> uctx
+  let nf_evar c = EConstr.to_constr evd (EConstr.of_constr c) in
+  let uvars, l = List.fold_left_map (fun uvars (coe,t,imps) ->
+      let t = nf_evar t in
+      let uvars = Univ.LSet.union uvars (Univops.universes_of_constr t) in
+      uvars, (coe,t,imps))
+      Univ.LSet.empty l
   in
-  let (_, _, st) = declare_assumption coe kind (ty, uctx) binders impls false nl id in
-  st
-
-let do_assumptions kind nl l = match l with
-| [coe, ([id, Some pl], c)] ->
-  let () = match kind with
-  | (Discharge, _, _) when Lib.sections_are_opened () ->
-    let loc = fst id in
-    let msg = Pp.str "Section variables cannot be polymorphic." in
-    user_err ?loc  msg
-  | _ -> ()
-  in
-  do_assumptions_bound_univs coe kind nl id (Some pl) c
-| _ ->
-  let map (coe, (idl, c)) =
-    let map (id, univs) = match univs with
-    | None -> id
-    | Some _ ->
-      let loc = fst id in
-      let msg =
-	Pp.str "Assumptions with bound universes can only be defined one at a time." in
-      user_err ?loc  msg
-    in
-    (coe, (List.map map idl, c))
-  in
-  let l = List.map map l in
-  do_assumptions_unbound_univs kind nl l
+  let evd = Evd.restrict_universe_context evd uvars in
+  let uctx = Evd.check_univ_decl ~poly:(pi2 kind) evd udecl in
+  let ubinders = Evd.universe_binders evd in
+  pi2 (List.fold_left (fun (subst,status,uctx) ((is_coe,idl),t,imps) ->
+      let t = replace_vars subst t in
+      let refs, status' = declare_assumptions idl is_coe kind (t,uctx) ubinders imps nl in
+      let subst' = List.map2
+          (fun (_,id) (c,u) -> (id, Universes.constr_of_global_univ (c,u)))
+          idl refs
+      in
+      subst'@subst, status' && status, next_uctx uctx)
+    ([], true, uctx) l)
 
 (* 3a| Elimination schemes for mutual inductive definitions *)
 
diff --git a/vernac/command.mli b/vernac/command.mli
index fb99a717b..070f3e112 100644
--- a/vernac/command.mli
+++ b/vernac/command.mli
@@ -43,7 +43,7 @@ val do_definition : Id.t -> definition_kind -> Vernacexpr.universe_decl_expr opt
 (** returns [false] if the assumption is neither local to a section,
     nor in a module type and meant to be instantiated. *)
 val declare_assumption : coercion_flag -> assumption_kind -> 
-  types Univ.in_universe_context_set ->
+  types in_constant_universes_entry ->
   Universes.universe_binders -> Impargs.manual_implicits ->
   bool (** implicit *) -> Vernacexpr.inline -> variable Loc.located ->
   global_reference * Univ.Instance.t * bool