EConstr iterators respect the binding structure of cases.

Fixes #3166.

3 files changed, 42 insertions, 16 deletions

diff --git a/engine/eConstr.ml b/engine/eConstr.ml
index ec255aad03..7748df1a9b 100644
--- a/engine/eConstr.ml
+++ b/engine/eConstr.ml
@@ -373,6 +373,18 @@ let expand_case env _sigma (ci, u, pms, p, iv, c, bl) =
   let bl = of_constr_array bl in
   (ci, p, iv, c, bl)
 
+let annotate_case env sigma (ci, u, pms, p, iv, c, bl as case) =
+  let (_, p, _, _, bl) = expand_case env sigma case in
+  let p =
+    (* Too bad we need to fetch this data in the environment, should be in the
+      case_info instead. *)
+    let (_, mip) = Inductive.lookup_mind_specif env ci.ci_ind in
+    decompose_lam_n_decls sigma (mip.Declarations.mind_nrealdecls + 1) p
+  in
+  let mk_br c n = decompose_lam_n_decls sigma n c in
+  let bl = Array.map2 mk_br bl ci.ci_cstr_ndecls in
+  (ci, u, pms, p, iv, c, bl)
+
 let expand_branch env _sigma u pms (ind, i) (nas, _br) =
   let open Declarations in
   let u = EInstance.unsafe_to_instance u in
@@ -413,9 +425,9 @@ let iter_with_full_binders env sigma g f n c =
   | App (c,l) -> f n c; Array.Fun1.iter f n l
   | Evar (_,l) -> List.iter (fun c -> f n c) l
   | Case (ci,u,pms,p,iv,c,bl) ->
-    (* FIXME: skip the type annotations *)
-    let (ci, p, iv, c, bl) = expand_case env sigma (ci, u, pms, p, iv, c, bl) in
-    f n p; iter_invert (f n) iv; f n c; Array.Fun1.iter f n bl
+    let (ci, _, pms, p, iv, c, bl) = annotate_case env sigma (ci, u, pms, p, iv, c, bl) in
+    let f_ctx (ctx, c) = f (List.fold_right g ctx n) c in
+    Array.Fun1.iter f n pms; f_ctx p; iter_invert (f n) iv; f n c; Array.iter f_ctx bl
   | Proj (p,c) -> f n c
   | Fix (_,(lna,tl,bl)) ->
     Array.iter (f n) tl;
diff --git a/engine/eConstr.mli b/engine/eConstr.mli
index 2add7b0b5e..e914d4f884 100644
--- a/engine/eConstr.mli
+++ b/engine/eConstr.mli
@@ -343,6 +343,10 @@ val is_global : Evd.evar_map -> GlobRef.t -> t -> bool
 val expand_case : Environ.env -> Evd.evar_map ->
   case -> (case_info * t * case_invert * t * t array)
 
+val annotate_case : Environ.env -> Evd.evar_map -> case ->
+  case_info * EInstance.t * t array * (rel_context * t) * case_invert * t * (rel_context * t) array
+(** Same as above, but doesn't turn contexts into binders *)
+
 val expand_branch : Environ.env -> Evd.evar_map ->
   EInstance.t -> t array -> constructor -> case_branch -> rel_context
 (** Given a universe instance and parameters for the inductive type,
diff --git a/engine/termops.ml b/engine/termops.ml
index d331ecc458..4dc584cfa8 100644
--- a/engine/termops.ml
+++ b/engine/termops.ml
@@ -651,15 +651,19 @@ let map_constr_with_binders_left_to_right env sigma g f l c =
       if List.for_all2 (==) al' al then c
       else mkEvar (e, al')
   | Case (ci,u,pms,p,iv,b,bl) ->
-      (* FIXME: skip annotations? *)
-      let (ci, p, iv, b, bl) = EConstr.expand_case env sigma (ci, u, pms, p, iv, b, bl) in
+      let (ci, _, pms, p0, _, b, bl0) = annotate_case env sigma (ci, u, pms, p, iv, b, bl) in
+      let f_ctx (nas, _ as r) (ctx, c) =
+        let c' = f (List.fold_right g ctx l) c in
+        if c' == c then r else (nas, c')
+      in
       (* In v8 concrete syntax, predicate is after the term to match! *)
       let b' = f l b in
+      let pms' = Array.map_left (f l) pms in
+      let p' = f_ctx p p0 in
       let iv' = map_invert (f l) iv in
-      let p' = f l p in
-      let bl' = Array.map_left (f l) bl in
-      if b' == b && p' == p && iv' == iv && bl' == bl then c
-      else mkCase (EConstr.contract_case env sigma (ci, p', iv', b', bl'))
+      let bl' = Array.map_left (fun (c, c0) -> f_ctx c c0) (Array.map2 (fun x y -> (x, y)) bl bl0) in
+        if b' == b && pms' == pms && p' == p && iv' == iv && bl' == bl then c
+        else mkCase (ci, u, pms', p', iv', b', bl')
   | Fix (ln,(lna,tl,bl as fx)) ->
       let l' = fold_rec_types g fx l in
       let (tl', bl') = map_left2 (f l) tl (f l') bl in
@@ -713,13 +717,18 @@ let map_constr_with_full_binders env sigma g f l cstr =
       let al' = List.map (f l) al in
       if List.for_all2 (==) al al' then cstr else mkEvar (e, al')
   | Case (ci, u, pms, p, iv, c, bl) ->
-      let (ci, p, iv, c, bl) = EConstr.expand_case env sigma (ci, u, pms, p, iv, c, bl) in
-      let p' = f l p in
+      let (ci, _, pms, p0, _, c, bl0) = annotate_case env sigma (ci, u, pms, p, iv, c, bl) in
+      let f_ctx (nas, _ as r) (ctx, c) =
+        let c' = f (List.fold_right g ctx l) c in
+        if c' == c then r else (nas, c')
+      in
+      let pms' = Array.Smart.map (f l) pms in
+      let p' = f_ctx p p0 in
       let iv' = map_invert (f l) iv in
       let c' = f l c in
-      let bl' = Array.map (f l) bl in
-      if p==p' && iv'==iv && c==c' && Array.for_all2 (==) bl bl' then cstr else
-        mkCase (EConstr.contract_case env sigma (ci, p', iv', c', bl'))
+      let bl' = Array.map2 f_ctx bl bl0 in
+      if pms==pms' && p==p' && iv'==iv && c==c' && Array.for_all2 (==) bl bl' then cstr else
+        mkCase (ci, u, pms', p', iv', c', bl')
   | Fix (ln,(lna,tl,bl as fx)) ->
       let tl' = Array.map (f l) tl in
       let l' = fold_rec_types g fx l in
@@ -760,8 +769,9 @@ let fold_constr_with_full_binders env sigma g f n acc c =
   | Proj (_,c) -> f n acc c
   | Evar (_,l) -> List.fold_left (f n) acc l
   | Case (ci, u, pms, p, iv, c, bl) ->
-    let (ci, p, iv, b, bl) = EConstr.expand_case env sigma (ci, u, pms, p, iv, c, bl) in
-    Array.fold_left (f n) (f n (fold_invert (f n) (f n acc p) iv) c) bl
+    let (ci, _, pms, p, _, c, bl) = EConstr.annotate_case env sigma (ci, u, pms, p, iv, c, bl) in
+    let f_ctx acc (ctx, c) = f (List.fold_right g ctx n) acc c in
+    Array.fold_left f_ctx (f n (fold_invert (f n) (f_ctx (Array.fold_left (f n) acc pms) p) iv) 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