Allow unfolding of mapi_with_cps to specialize the function to the type argument it gets passed

1 files changed, 33 insertions, 10 deletions

diff --git a/src/Util/CPSUtil.v b/src/Util/CPSUtil.v
index e75eda04d..7157867ba 100644
--- a/src/Util/CPSUtil.v
+++ b/src/Util/CPSUtil.v
@@ -283,18 +283,29 @@ Module Tuple.
     [|rewrite IHn, <-map2_append,<-!subst_append]; reflexivity.
   Qed. Hint Rewrite @map2_cps_correct : uncps.
 
-  Fixpoint mapi_with'_cps {T A B n} i
+  Section internal_mapi_with_cps.
+    (* We define fixpoints with fewer parameters to the internal [fix] to allow unfolding to partially specialize them *)
+    Context {R T A B : Type}
+            (f: nat->T->A->(T*B->R)->R).
+
+    Fixpoint mapi_with'_cps_specialized {n} i
+             (start:T)
+      : Tuple.tuple' A n -> (T * tuple' B n -> R) -> R :=
+      match n as n0 return (tuple' A n0 -> (T * tuple' B n0->R)->R) with
+      | O => fun ys ret => f i start ys ret
+      | S n' => fun ys ret =>
+                  f i start (hd ys) (fun sb =>
+                  mapi_with'_cps_specialized (S i) (fst sb) (tl ys)
+                      (fun r => ret (fst r, (snd r, snd sb))))
+      end.
+  End internal_mapi_with_cps.
+
+  Definition mapi_with'_cps {T A B n} i
           (f: nat->T->A->forall {R}, (T*B->R)->R) (start:T)
-  : Tuple.tuple' A n -> forall {R}, (T * tuple' B n -> R) -> R :=
-  match n as n0 return (tuple' A n0 -> forall {R}, (T * tuple' B n0->R)->R) with
-  | O => fun ys {T} ret => f i start ys ret
-  | S n' => fun ys {T} ret =>
-              f i start (hd ys) (fun sb =>
-              mapi_with'_cps (S i) f (fst sb) (tl ys)
-                  (fun r => ret (fst r, (snd r, snd sb))))
-  end.
+    : Tuple.tuple' A n -> forall {R}, (T * tuple' B n -> R) -> R
+    := fun ts R => @mapi_with'_cps_specialized R T A B (fun n t a => @f n t a R) n i start ts.
 
-  Fixpoint mapi_with_cps {S A B n}
+  Definition mapi_with_cps {S A B n}
           (f: nat->S->A->forall {T}, (S*B->T)->T) (start:S)
   : tuple A n -> forall {T}, (S * tuple B n->T)->T :=
   match n as n0 return (tuple A n0 -> forall {T}, (S * tuple B n0->T)->T) with
@@ -302,6 +313,18 @@ Module Tuple.
   | S n' => fun ys {T} ret => mapi_with'_cps 0%nat f start ys ret
   end.
 
+  Lemma unfold_mapi_with'_cps {T A B n} i
+        (f: nat->T->A->forall {R}, (T*B->R)->R) (start:T)
+    : @mapi_with'_cps T A B n i f start
+      = match n as n0 return (tuple' A n0 -> forall {R}, (T * tuple' B n0->R)->R) with
+        | O => fun ys {T} ret => f i start ys ret
+        | S n' => fun ys {T} ret =>
+                    f i start (hd ys) (fun sb =>
+                    mapi_with'_cps (S i) f (fst sb) (tl ys)
+                        (fun r => ret (fst r, (snd r, snd sb))))
+        end.
+  Proof. destruct n; reflexivity. Qed.
+
   Lemma mapi_with'_cps_correct {S A B n} : forall i f start xs T ret,
   (forall i s a R (ret:_->R), f i s a R ret = ret (f i s a _ id)) ->
   @mapi_with'_cps S A B n i f start xs T ret = ret (mapi_with' i (fun i s a => f i s a _ id) start xs).