Fix lifting in foo_with_full_binders for (co)fixpoints

5 files changed, 82 insertions, 51 deletions

diff --git a/engine/eConstr.ml b/engine/eConstr.ml
index cfc4bea85f..40e48998d6 100644
--- a/engine/eConstr.ml
+++ b/engine/eConstr.ml
@@ -296,6 +296,8 @@ let decompose_prod_n_assum sigma n c =
 
 let existential_type = Evd.existential_type
 
+let lift n c = of_constr (Vars.lift n (unsafe_to_constr c))
+
 let map_under_context f n c =
   let f c = unsafe_to_constr (f (of_constr c)) in
   of_constr (Constr.map_under_context f n (unsafe_to_constr c))
@@ -453,11 +455,11 @@ let iter_with_full_binders sigma g f n c =
   | Proj (p,c) -> f n c
   | Fix (_,(lna,tl,bl)) ->
     Array.iter (f n) tl;
-    let n' = Array.fold_left2 (fun n na t -> g (LocalAssum (na,t)) n) n lna tl in
+    let n' = Array.fold_left2_i (fun i n na t -> g (LocalAssum (na, lift i t)) n) n lna tl in
     Array.iter (f n') bl
   | CoFix (_,(lna,tl,bl)) ->
     Array.iter (f n) tl;
-    let n' = Array.fold_left2 (fun n na t -> g (LocalAssum (na,t)) n) n lna tl in
+    let n' = Array.fold_left2_i (fun i n na t -> g (LocalAssum (na,lift i t)) n) n lna tl in
     Array.iter (f n') bl
 
 let iter_with_binders sigma g f n c =
@@ -712,7 +714,7 @@ let to_rel_decl = unsafe_to_rel_decl
 type substl = t list
 
 (** Operations that commute with evar-normalization *)
-let lift n c = of_constr (Vars.lift n (to_constr c))
+let lift = lift
 let liftn n m c = of_constr (Vars.liftn n m (to_constr c))
 
 let substnl subst n c = of_constr (Vars.substnl (cast_list unsafe_eq subst) n (to_constr c))
diff --git a/engine/termops.ml b/engine/termops.ml
index ada6311067..64048777ad 100644
--- a/engine/termops.ml
+++ b/engine/termops.ml
@@ -721,18 +721,16 @@ let map_constr_with_full_binders_gen userview sigma g f l cstr =
       let bl' = Array.map (f l) bl in
       if p==p' && c==c' && Array.for_all2 (==) bl bl' then cstr else
         mkCase (ci, p', c', bl')
-  | Fix (ln,(lna,tl,bl)) ->
+  | Fix (ln,(lna,tl,bl as fx)) ->
       let tl' = Array.map (f l) tl in
-      let l' =
-        Array.fold_left2 (fun l na t -> g (RelDecl.LocalAssum (na, t)) l) l lna tl in
+      let l' = fold_rec_types g fx l in
       let bl' = Array.map (f l') bl in
       if Array.for_all2 (==) tl tl' && Array.for_all2 (==) bl bl'
       then cstr
       else mkFix (ln,(lna,tl',bl'))
-  | CoFix(ln,(lna,tl,bl)) ->
+  | CoFix(ln,(lna,tl,bl as fx)) ->
       let tl' = Array.map (f l) tl in
-      let l' =
-        Array.fold_left2 (fun l na t -> g (RelDecl.LocalAssum (na, t)) l) l lna tl in
+      let l' = fold_rec_types g fx l in
       let bl' = Array.map (f l') bl in
       if Array.for_all2 (==) tl tl' && Array.for_all2 (==) bl bl'
       then cstr
@@ -759,7 +757,10 @@ let fold_constr_with_full_binders sigma g f n acc c =
   Constr.fold_with_full_binders g f n acc c
 
 let fold_constr_with_binders sigma g f n acc c =
-  fold_constr_with_full_binders sigma (fun _ x -> g x) f n acc c
+  let open EConstr in
+  let f l acc c = f l acc (of_constr c) in
+  let c = Unsafe.to_constr (whd_evar sigma c) in
+  Constr.fold_constr_with_binders g f n acc c
 
 (* [iter_constr_with_full_binders g f acc c] iters [f acc] on the immediate
    subterms of [c]; it carries an extra data [acc] which is processed by [g] at
@@ -780,13 +781,19 @@ let iter_constr_with_full_binders sigma g f l c =
   | Evar (_,args) -> Array.iter (f l) args
   | Case (_,p,c,bl) -> f l p; f l c; Array.iter (f l) bl
   | Fix (_,(lna,tl,bl)) ->
-      let l' = Array.fold_left2 (fun l na t -> g (LocalAssum (na,t)) l) l lna tl in
-      Array.iter (f l) tl;
-      Array.iter (f l') bl
+    let l' = Array.fold_left2_i (fun i l na t ->
+        g (LocalAssum (na, EConstr.Vars.lift i t)) l)
+        l lna tl
+    in
+    Array.iter (f l) tl;
+    Array.iter (f l') bl
   | CoFix (_,(lna,tl,bl)) ->
-      let l' = Array.fold_left2 (fun l na t -> g (LocalAssum (na,t)) l) l lna tl in
-      Array.iter (f l) tl;
-      Array.iter (f l') bl
+    let l' = Array.fold_left2_i (fun i l na t ->
+        g (LocalAssum (na, EConstr.Vars.lift i t)) l)
+        l lna tl
+    in
+    Array.iter (f l) tl;
+    Array.iter (f l') bl
 
 (***************************)
 (* occurs check functions  *)
diff --git a/kernel/constr.ml b/kernel/constr.ml
index b60a2a4b46..8e5d15dd2d 100644
--- a/kernel/constr.ml
+++ b/kernel/constr.ml
@@ -452,27 +452,6 @@ let fold f acc c = match kind c with
   | CoFix (_,(_lna,tl,bl)) ->
     Array.fold_left2 (fun acc t b -> f (f acc t) b) acc tl bl
 
-let fold_with_full_binders g f n acc c =
-  let open Context.Rel.Declaration in
-  match kind c with
-  | Rel _ | Meta _ | Var _   | Sort _ | Const _ | Ind _ | Construct _ -> acc
-  | Cast (c,_, t) -> f n (f n acc c) t
-  | Prod (na,t,c) -> f (g (LocalAssum (na,t)) n) (f n acc t) c
-  | Lambda (na,t,c) -> f (g (LocalAssum (na,t)) n) (f n acc t) c
-  | LetIn (na,b,t,c) -> f (g (LocalDef (na,b,t)) n) (f n (f n acc b) t) c
-  | App (c,l) -> Array.fold_left (f n) (f n acc c) l
-  | Proj (_,c) -> f n acc c
-  | Evar (_,l) -> Array.fold_left (f n) acc l
-  | Case (_,p,c,bl) -> Array.fold_left (f n) (f n (f n acc p) c) bl
-  | Fix (_,(lna,tl,bl)) ->
-      let n' = CArray.fold_left2 (fun c n t -> g (LocalAssum (n,t)) c) n lna tl in
-      let fd = Array.map2 (fun t b -> (t,b)) tl bl in
-      Array.fold_left (fun acc (t,b) -> f n' (f n acc t) b) acc fd
-  | CoFix (_,(lna,tl,bl)) ->
-      let n' = CArray.fold_left2 (fun c n t -> g (LocalAssum (n,t)) c) n lna tl in
-      let fd = Array.map2 (fun t b -> (t,b)) tl bl in
-      Array.fold_left (fun acc (t,b) -> f n' (f n acc t) b) acc fd
-
 (* [iter f c] iters [f] on the immediate subterms of [c]; it is
    not recursive and the order with which subterms are processed is
    not specified *)
@@ -799,6 +778,49 @@ let map_with_binders g f l c0 = match kind c0 with
     let bl' = Array.Fun1.Smart.map f l' bl in
     mkCoFix (ln,(lna,tl',bl'))
 
+(*********************)
+(*      Lifting      *)
+(*********************)
+
+(* The generic lifting function *)
+let rec exliftn el c =
+  let open Esubst in
+  match kind c with
+  | Rel i -> mkRel(reloc_rel i el)
+  | _ -> map_with_binders el_lift exliftn el c
+
+(* Lifting the binding depth across k bindings *)
+
+let liftn n k c =
+  let open Esubst in
+  match el_liftn (pred k) (el_shft n el_id) with
+    | ELID -> c
+    | el -> exliftn el c
+
+let lift n = liftn n 1
+
+let fold_with_full_binders g f n acc c =
+  let open Context.Rel.Declaration in
+  match kind c with
+  | Rel _ | Meta _ | Var _   | Sort _ | Const _ | Ind _ | Construct _ -> acc
+  | Cast (c,_, t) -> f n (f n acc c) t
+  | Prod (na,t,c) -> f (g (LocalAssum (na,t)) n) (f n acc t) c
+  | Lambda (na,t,c) -> f (g (LocalAssum (na,t)) n) (f n acc t) c
+  | LetIn (na,b,t,c) -> f (g (LocalDef (na,b,t)) n) (f n (f n acc b) t) c
+  | App (c,l) -> Array.fold_left (f n) (f n acc c) l
+  | Proj (_,c) -> f n acc c
+  | Evar (_,l) -> Array.fold_left (f n) acc l
+  | Case (_,p,c,bl) -> Array.fold_left (f n) (f n (f n acc p) c) bl
+  | Fix (_,(lna,tl,bl)) ->
+      let n' = CArray.fold_left2_i (fun i c n t -> g (LocalAssum (n,lift i t)) c) n lna tl in
+      let fd = Array.map2 (fun t b -> (t,b)) tl bl in
+      Array.fold_left (fun acc (t,b) -> f n' (f n acc t) b) acc fd
+  | CoFix (_,(lna,tl,bl)) ->
+      let n' = CArray.fold_left2_i (fun i c n t -> g (LocalAssum (n,lift i t)) c) n lna tl in
+      let fd = Array.map2 (fun t b -> (t,b)) tl bl in
+      Array.fold_left (fun acc (t,b) -> f n' (f n acc t) b) acc fd
+
+
 type 'univs instance_compare_fn = GlobRef.t -> int ->
   'univs -> 'univs -> bool
 
diff --git a/kernel/constr.mli b/kernel/constr.mli
index 1be1f63ff7..f2cedcdabb 100644
--- a/kernel/constr.mli
+++ b/kernel/constr.mli
@@ -383,6 +383,17 @@ type rel_context = rel_declaration list
 type named_context = named_declaration list
 type compacted_context = compacted_declaration list
 
+(** {6 Relocation and substitution } *)
+
+(** [exliftn el c] lifts [c] with lifting [el] *)
+val exliftn : Esubst.lift -> constr -> constr
+
+(** [liftn n k c] lifts by [n] indexes above or equal to [k] in [c] *)
+val liftn : int -> int -> constr -> constr
+
+(** [lift n c] lifts by [n] the positive indexes in [c] *)
+val lift : int -> constr -> constr
+
 (** {6 Functionals working on expressions canonically abstracted over
        a local context (possibly with let-ins)} *)
 
diff --git a/kernel/vars.ml b/kernel/vars.ml
index 7380a860dd..f9c576ca4a 100644
--- a/kernel/vars.ml
+++ b/kernel/vars.ml
@@ -9,7 +9,6 @@
 (************************************************************************)
 
 open Names
-open Esubst
 
 module RelDecl = Context.Rel.Declaration
 
@@ -80,19 +79,9 @@ let noccur_with_meta n m term =
 (*      Lifting      *)
 (*********************)
 
-(* The generic lifting function *)
-let rec exliftn el c = match Constr.kind c with
-  | Constr.Rel i -> Constr.mkRel(reloc_rel i el)
-  | _ -> Constr.map_with_binders el_lift exliftn el c
-
-(* Lifting the binding depth across k bindings *)
-
-let liftn n k c =
-  match el_liftn (pred k) (el_shft n el_id) with
-    | ELID -> c
-    | el -> exliftn el c
-
-let lift n = liftn n 1
+let exliftn = Constr.exliftn
+let liftn = Constr.liftn
+let lift = Constr.lift
 
 (*********************)
 (*   Substituting    *)