Add fresh annotations to Consts.

kernel/constr.ml

@@ -816,6 +816,7 @@ let rec count_annots cstr =
     count_annots c
   | LetIn (_, b, _, c) ->
     count_annots b + count_annots c
+  | Const (_, la) -> List.length @@ Option.default [] la
   | Case (_, _, c, lf) ->
     Array.fold_left (fun count c -> count + count_annots c) (count_annots c) lf
   | Fix (_, (_, _, bl))
@@ -836,10 +837,22 @@ let rec any_annot f c =
 
 (** map-type functions on stage annotations of constrs *)
 
+let map_annots_list g cstr =
+  match cstr with
+  | Rel (n, Some la) ->
+    let la' = g la in
+    if la == la' then cstr else
+    mkRelA n (Some la')
+  | Const (c, Some la) ->
+    let la' = g la in
+    if la == la' then cstr else
+    mkConstUA c (Some la')
+  | _ -> cstr
+
 let rec map_annots f g cstr =
   match cstr with
   | Ind (iu, a) -> f iu a cstr
-  | Rel (n, la) -> g n la cstr
+  | Rel _ -> map_annots_list g cstr
   | Cast (c, k, t) ->
     let c' = map_annots f g c in
     if c == c' then cstr else
@@ -853,6 +866,7 @@ let rec map_annots f g cstr =
     let c' = map_annots f g c in
     if b == b' && c == c' then cstr else
     mkLetIn (n, b', t, c')
+  | Const _ -> map_annots_list g cstr
   | Case (ci, p, c, lf) ->
     let c' = map_annots f g c in
     let lf' = Array.Smart.map (map_annots f g) lf in
@@ -868,26 +882,22 @@ let rec map_annots f g cstr =
     mkCoFix (ln, (nl, tl, bl'))
   | _ -> map (map_annots f g) cstr
 
-let mkRelMap annot n la c =
-  match la with
-  | None -> c
-  | Some la ->
-    mkRelA n @@ Some (List.make (List.length la) annot)
-let const2 _ _ = identity
+let mkListAnnot annot la =
+  List.make (List.length la) annot
 
 let erase =
   let f iu a c =
     match a with
     | Empty -> c
     | _ -> mkIndUS iu Empty in
-  map_annots f @@ mkRelMap Empty
+  map_annots f @@ mkListAnnot Empty
 
 let erase_infty =
   let f iu a c =
     match a with
     | Stage Infty -> c
     | _ -> mkIndUS iu infty in
-  map_annots f @@ mkRelMap infty
+  map_annots f @@ mkListAnnot infty
 
 let erase_glob vars =
   let f iu a c =
@@ -897,7 +907,7 @@ let erase_glob vars =
       mkIndUS iu Glob
     | Stage Infty -> c
     | _ -> mkIndUS iu infty in
-  map_annots f @@ mkRelMap infty
+  map_annots f @@ mkListAnnot infty
 
 let erase_star vars =
   let f iu a c =
@@ -907,37 +917,33 @@ let erase_star vars =
       mkIndUS iu Star
     | Empty -> c
     | _ -> mkIndUS iu Empty in
-  map_annots f @@ mkRelMap Empty
+  map_annots f @@ mkListAnnot Empty
 
 let annotate ind s =
   let f (((i, _), _) as iu) _ c =
     if MutInd.equal ind i then
       mkIndUS iu s
     else c in
-  map_annots f const2
+  map_annots f identity
 
 let annotate_fresh annots =
   let annots_ref = ref annots in
   let f iu _ _ =
     let annots = !annots_ref in
     annots_ref := List.tl annots;
     mkIndUS iu (List.hd annots) in
-  let g n la c =
-    match la with
-    | None -> c
-    | Some la ->
-      let annots = !annots_ref in
-      let la, annots = List.chop (List.length la) annots in
-      annots_ref := annots;
-      mkRelA n (Some la) in
+  let g la =
+    let annots = !annots_ref in
+    let la, annots = List.chop (List.length la) annots in
+    annots_ref := annots; la in
   map_annots f g
 
 let annotate_glob s =
   let f iu a c =
     match a with
     | Glob -> mkIndUS iu s
     | _ -> c in
-  map_annots f const2
+  map_annots f identity
 
 let annotate_succ vars =
   let f iu a c =
@@ -946,7 +952,7 @@ let annotate_succ vars =
       when mem na vars ->
       mkIndUS iu (hat a)
     | _ -> c in
-  map_annots f const2
+  map_annots f identity
 
 (*********************)
 (*      Lifting      *)

kernel/typeops.ml

@@ -170,6 +170,12 @@ let type_of_constant_in env (kn,_u as cst) =
   let cstrnts = check_hyps_inclusion env mkConstU cst cb.const_hyps in
   constant_type_in env cst, cstrnts
 
+let stage_vars_in_constant env (kn, _) =
+  let cb = lookup_constant kn env in
+  match cb.const_body with
+  | Def c -> count_annots @@ Mod_subst.force_constr c
+  | _ -> 0
+
 (* Type of a lambda-abstraction. *)
 
 (* [judge_of_abstraction env name var j] implements the rule
@@ -548,10 +554,12 @@ let rec execute env stg cstr =
       stg, empty (), cstr, type_of_variable env id
 
     | Const (c, _ans) ->
+      let numvars = stage_vars_in_constant env c in
       let t, cstrnt = type_of_constant env c in
       let s, stg = next stg in
+      let annots, stg = next_annots numvars stg in
       let t = annotate_glob s t in
-      stg, cstrnt, cstr, t
+      stg, cstrnt, mkConstUA c (Some annots), t
 
     | Proj (p, c) ->
       let stg, cstrnt, c', ct = execute env stg c in