From Coq Require Import List. Import ListNotations.
From Coq Require Import ZArith Lia.
Require Import Crypto.Util.Tactics.BreakMatch.
Require Import Crypto.Util.Tactics.DestructHead.
Require Import Crypto.Util.ListUtil.
Require Import Crypto.Util.IdfunWithAlt.
Require Import Crypto.Util.LetIn.
Require Crypto.Util.Tuple. Local Notation tuple := Tuple.tuple.
Local Open Scope Z_scope.

Create HintDb push_id discriminated.
Create HintDb uncps discriminated.

Lemma push_id {A} (a:A) : id a = a. reflexivity. Qed.
#[global]
Hint Rewrite @push_id : push_id.

Definition id_with_alt_cps' {R A} (value value_for_alt : (A -> A) -> A) (f : A -> R) : R
  := dlet z := id_with_alt (value id) (value_for_alt id) in
     f z.
Definition id_with_alt_cps'' {R A} (value value_for_alt : (A -> R) -> R) (f : A -> R) : R
  := id_with_alt (value f) (value_for_alt f).
Definition id_with_alt_cps {R A} (value value_for_alt : forall R, (A -> R) -> R) (f : A -> R) : R
  := id_with_alt_cps' (value _) (value_for_alt _) f.
Definition id_with_alt_cps'_correct {R A} value value_for_alt f
  : @id_with_alt_cps' R A value value_for_alt f = f (id_with_alt (value id) (value_for_alt id))
  := eq_refl.
#[global]
Hint Rewrite @id_with_alt_cps'_correct : uncps.
Definition id_with_alt_cps''_correct {R A} value value_for_alt f
  : @id_with_alt_cps'' R A value value_for_alt f = id_with_alt (value f) (value_for_alt f)
  := eq_refl.
#[global]
Hint Rewrite @id_with_alt_cps''_correct : uncps.
Definition id_with_alt_cps_correct {R A} value value_for_alt f
  : @id_with_alt_cps R A value value_for_alt f = f (id_with_alt (value _ id) (value_for_alt _ id))
  := eq_refl.
#[global]
Hint Rewrite @id_with_alt_cps_correct : uncps.
Definition id_with_alt_cps'_correct_gen {R A} (value value_for_alt : forall R, (A -> R) -> R) f
  : @id_with_alt_cps' R A (value _) (value_for_alt _) f = f (id_with_alt (value _ id) (value_for_alt _ id))
  := eq_refl.
#[global]
Hint Rewrite @id_with_alt_cps'_correct_gen : uncps.
Definition id_with_alt_cps''_correct_gen {R A} (value value_for_alt : forall R, (A -> R) -> R) f
           (Hvalue : value _ f = f (value _ id))
  : @id_with_alt_cps'' R A (value _) (value_for_alt _) f = f (id_with_alt (value _ id) (value_for_alt _ id)).
Proof.
  cbv [id_with_alt_cps'' id_with_alt]; assumption.
Defined.
#[global]
Hint Rewrite @id_with_alt_cps''_correct_gen : uncps.

Definition id_tuple'_with_alt_cps' {R A n}
           (value value_for_alt : (Tuple.tuple' A n -> Tuple.tuple' A n) -> Tuple.tuple' A n)
           (f : Tuple.tuple' A n -> R)
  : R
  := dlet z := id_tuple'_with_alt (value id) (value_for_alt id) in
     f z.
Definition id_tuple_with_alt_cps' {R A n}
           (value value_for_alt : (Tuple.tuple A n -> Tuple.tuple A n) -> Tuple.tuple A n)
           (f : Tuple.tuple A n -> R)
  : R
  := dlet z := id_tuple_with_alt (value id) (value_for_alt id) in
     f z.

Definition id_tuple'_with_alt_cps'_correct {R A n} value value_for_alt f
  : @id_tuple'_with_alt_cps' R A n value value_for_alt f = f (id_tuple'_with_alt (value id) (value_for_alt id))
  := eq_refl.
#[global]
Hint Rewrite @id_tuple'_with_alt_cps'_correct : uncps.

Definition id_tuple_with_alt_cps'_correct {R A n} value value_for_alt f
  : @id_tuple_with_alt_cps' R A n value value_for_alt f = f (id_tuple_with_alt (value id) (value_for_alt id))
  := eq_refl.
#[global]
Hint Rewrite @id_tuple_with_alt_cps'_correct : uncps.

Definition id_tuple'_with_alt_cps'' {R A n}
           (value value_for_alt : (Tuple.tuple' A n -> R) -> R)
           (f : Tuple.tuple' A n -> R)
  : R
  := id_with_alt (value f) (value_for_alt f).
Definition id_tuple_with_alt_cps'' {R A n}
           (value value_for_alt : (Tuple.tuple A n -> R) -> R)
           (f : Tuple.tuple A n -> R)
  : R
  := id_with_alt (value f) (value_for_alt f).

Definition id_tuple'_with_alt_cps''_correct {R A n} value value_for_alt f
  : @id_tuple'_with_alt_cps'' R A n value value_for_alt f = id_with_alt (value f) (value_for_alt f)
  := eq_refl.
#[global]
Hint Rewrite @id_tuple'_with_alt_cps''_correct : uncps.

Definition id_tuple_with_alt_cps''_correct {R A n} value value_for_alt f
  : @id_tuple_with_alt_cps'' R A n value value_for_alt f = id_with_alt (value f) (value_for_alt f)
  := eq_refl.
#[global]
Hint Rewrite @id_tuple_with_alt_cps''_correct : uncps.

Definition id_tuple'_with_alt_cps''_correct_gen {R A n}
           (value value_for_alt : forall R, (_ -> R) -> R) f
           (Hvalue : value _ f = f (value _ id))
  : @id_tuple'_with_alt_cps'' R A n (value _) (value_for_alt _) f = f (id_tuple'_with_alt (value _ id) (value_for_alt _ id)).
Proof.
  autorewrite with uncps.
  rewrite ?unfold_id_tuple'_with_alt, ?unfold_id_tuple_with_alt, ?unfold_id_with_alt, Hvalue.
  reflexivity.
Qed.
#[global]
Hint Rewrite @id_tuple'_with_alt_cps''_correct_gen : uncps.

Definition id_tuple_with_alt_cps''_correct_gen {R A n}
           (value value_for_alt : forall R, (_ -> R) -> R) f
           (Hvalue : value _ f = f (value _ id))
  : @id_tuple_with_alt_cps'' R A n (value _) (value_for_alt _) f = f (id_tuple_with_alt (value _ id) (value_for_alt _ id)).
Proof.
  autorewrite with uncps.
  rewrite ?unfold_id_tuple_with_alt, ?unfold_id_tuple_with_alt, ?unfold_id_with_alt, Hvalue.
  reflexivity.
Qed.
#[global]
Hint Rewrite @id_tuple_with_alt_cps''_correct_gen : uncps.

Definition id_tuple'_with_alt_cps {R A n}
           (value value_for_alt : forall R, (Tuple.tuple' A n -> R) -> R)
           (f : Tuple.tuple' A n -> R)
  : R
  := id_tuple'_with_alt_cps' (value _) (value_for_alt _) f.
Definition id_tuple_with_alt_cps {R A n}
           (value value_for_alt : forall R, (Tuple.tuple A n -> R) -> R)
           (f : Tuple.tuple A n -> R)
  : R
  := id_tuple_with_alt_cps' (value _) (value_for_alt _) f.

Definition id_tuple'_with_alt_cps_correct {R A n} value value_for_alt f
  : @id_tuple'_with_alt_cps R A n value value_for_alt f = f (id_tuple'_with_alt (value _ id) (value_for_alt _ id))
  := eq_refl.
#[global]
Hint Rewrite @id_tuple'_with_alt_cps_correct : uncps.

Definition id_tuple_with_alt_cps_correct {R A n} value value_for_alt f
  : @id_tuple_with_alt_cps R A n value value_for_alt f = f (id_tuple_with_alt (value _ id) (value_for_alt _ id))
  := eq_refl.
#[global]
Hint Rewrite @id_tuple_with_alt_cps_correct : uncps.

Lemma update_nth_id {T} i (xs:list T) : ListUtil.update_nth i id xs = xs.
Proof.
  revert xs; induction i; destruct xs; simpl; solve [ trivial | congruence ].
Qed.

Lemma map_fst_combine {A B} (xs:list A) (ys:list B) : List.map fst (List.combine xs ys) = List.firstn (length ys) xs.
Proof.
  revert xs; induction ys; destruct xs; simpl; solve [ trivial | congruence ].
Qed.

Lemma map_snd_combine {A B} (xs:list A) (ys:list B) : List.map snd (List.combine xs ys) = List.firstn (length xs) ys.
Proof.
  revert xs; induction ys; destruct xs; simpl; solve [ trivial | congruence ].
Qed.

Lemma nth_default_seq_inbouns d s n i (H:(i < n)%nat) :
  List.nth_default d (List.seq s n) i = (s+i)%nat.
Proof.
  progress cbv [List.nth_default].
  rewrite ListUtil.nth_error_seq.
  break_innermost_match; solve [ trivial | lia ].
Qed.

Lemma mod_add_mul_full a b c k m : m <> 0 -> c mod m = k mod m ->
                                   (a + b * c) mod m = (a + b * k) mod m.
Proof.
  intros; rewrite Z.add_mod, Z.mul_mod by auto.
  match goal with H : _ mod _ = _ mod _ |- _ => rewrite H end.
  rewrite <-Z.mul_mod, <-Z.add_mod by auto; reflexivity.
Qed.

Fixpoint map_cps {A B} (g : A->B) ls
         {T} (f:list B->T):=
  match ls with
  | nil => f nil
  | a :: t => map_cps g t (fun r => f (g a :: r))
  end.
Lemma map_cps_correct {A B} g ls: forall {T} f,
    @map_cps A B g ls T f = f (map g ls).
Proof. induction ls as [|?? IHls]; simpl; intros; rewrite ?IHls; reflexivity. Qed.
#[global]
Hint Rewrite @map_cps_correct : uncps.

Fixpoint map_cps2 {A B} (g : A->forall T, (B -> T) -> T) ls
         {T} (f:list B->T) : T:=
  match ls with
  | nil => f nil
  | a :: t => map_cps2 g t (fun r => g a _ (fun ga => f (ga :: r)))
  end.
Lemma map_cps2_correct {A B} g g' (Hg : forall T k a, g a T k = k (g' a)) ls : forall {T} f,
    @map_cps2 A B g ls T f = f (map g' ls).
Proof. induction ls as [|?? IHls]; simpl; intros; rewrite ?IHls, ?Hg; reflexivity. Qed.

Definition firstn_cps {A} (n:nat) (l:list A) {T} (f:list A->T) :=
  match n with
  | O => f nil
  | S n' => match l with
            | nil => f nil
            | a :: l' => f (a :: firstn n' l')
            end
  end.
Lemma firstn_cps_correct {A} n l T f :
  @firstn_cps A n l T f = f (firstn n l).
Proof. induction n; destruct l; reflexivity. Qed.
#[global]
Hint Rewrite @firstn_cps_correct : uncps.

Fixpoint flat_map_cps_specialized {T A B} (g:A->(list B->T)->T) (ls : list A) (f:list B->T)  :=
  match ls with
  | nil => f nil
  | (x::tl)%list => g x (fun r => flat_map_cps_specialized g tl (fun rr => f (r ++ rr))%list)
  end.

Definition flat_map_cps {A B} (g:A->forall {T}, (list B->T)->T) (ls : list A) {T} (f:list B->T)
  := @flat_map_cps_specialized T A B (fun a => @g a T) ls f.
Lemma unfold_flat_map_cps {A B} (g:A->forall {T}, (list B->T)->T) (ls : list A) {T} (f:list B->T)
  : @flat_map_cps A B g ls T f
    = match ls with
      | nil => f nil
      | (x::tl)%list => g x (fun r => flat_map_cps g tl (fun rr => f (r ++ rr))%list)
      end.
Proof. destruct ls; reflexivity. Qed.
Lemma flat_map_cps_correct {A B} (g:A->forall {T}, (list B->T)->T) ls :
  forall {T} (f:list B->T),
    (forall x T h, @g x T h = h (g x id)) ->
    @flat_map_cps A B g ls T f = f (List.flat_map (fun x => g x id) ls).
Proof.
  induction ls as [|?? IHls]; intros T f H; [reflexivity|].
  rewrite unfold_flat_map_cps. simpl @flat_map.
  rewrite H; erewrite IHls by eassumption.
  reflexivity.
Qed.
#[global]
Hint Rewrite @flat_map_cps_correct using (intros; autorewrite with uncps; auto): uncps.

Fixpoint from_list_default'_cps {A} (d y:A) n xs {T}:
  (Tuple.tuple' A n -> T) -> T:=
  match n as n0 return ((Tuple.tuple' A n0 ->T) ->T) with
  | O => fun f => f y
  | S n' => fun f =>
              match xs with
              | nil => from_list_default'_cps d d n' nil (fun r => f (r, y))
              | x :: xs' => from_list_default'_cps d x n' xs' (fun r => f (r, y))
              end
  end.
Lemma from_list_default'_cps_correct {A} n : forall d y l {T} f,
    @from_list_default'_cps A d y n l T f = f (Tuple.from_list_default' d y n l).
Proof.
  induction n as [|? IHn]; intros; simpl; [reflexivity|].
  break_match; subst; apply IHn.
Qed.
Definition from_list_default_cps {A} (d:A) n (xs:list A) {T} :
  (Tuple.tuple A n -> T) -> T:=
  match n as n0 return ((Tuple.tuple A n0 ->T) ->T) with
  | O => fun f => f tt
  | S n' => fun f =>
              match xs with
              | nil => from_list_default'_cps d d n' nil f
              | x :: xs' => from_list_default'_cps d x n' xs' f
              end
  end.
Lemma from_list_default_cps_correct {A} n : forall d l {T} f,
    @from_list_default_cps A d n l T f = f (Tuple.from_list_default d n l).
Proof.
  destruct n; intros; simpl; [reflexivity|].
  break_match; auto using from_list_default'_cps_correct.
Qed.
#[global]
Hint Rewrite @from_list_default_cps_correct : uncps.
Fixpoint to_list'_cps {A} n
         {T} (f:list A -> T) : Tuple.tuple' A n -> T :=
  match n as n0 return (Tuple.tuple' A n0 -> T) with
  | O => fun x => f [x]
  | S n' => fun (xs: Tuple.tuple' A (S n')) =>
              let (xs', x) := xs in
              to_list'_cps n' (fun r => f (x::r)) xs'
  end.
Lemma to_list'_cps_correct {A} n: forall t {T} f,
    @to_list'_cps A n T f t = f (Tuple.to_list' n t).
Proof.
  induction n; simpl; intros; [reflexivity|].
  destruct_head prod. apply IHn.
Qed.
Definition to_list_cps' {A} n {T} (f:list A->T)
  : Tuple.tuple A n -> T :=
  match n as n0 return (Tuple.tuple A n0 ->T) with
  | O => fun _ => f nil
  | S n' => to_list'_cps n' f
  end.
Definition to_list_cps {A} n t {T} f :=
  @to_list_cps' A n T f t.
Lemma to_list_cps_correct {A} n t {T} f :
  @to_list_cps A n t T f = f (Tuple.to_list n t).
Proof. cbv [to_list_cps to_list_cps' Tuple.to_list]; break_match; auto using to_list'_cps_correct. Qed.
#[global]
Hint Rewrite @to_list_cps_correct : uncps.

Definition on_tuple_cps {A B} (d:B) (g:list A ->forall {T},(list B->T)->T) {n m}
           (xs : Tuple.tuple A n) {T} (f:tuple B m ->T) :=
  to_list_cps n xs (fun r => g r (fun rr => from_list_default_cps d m rr f)).
Lemma on_tuple_cps_correct {A B} d (g:list A -> forall {T}, (list B->T)->T)
      {n m} xs {T} f
      (Hg : forall x {T} h, @g x T h = h (g x id)) : forall H,
    @on_tuple_cps A B d g n m xs T f = f (@Tuple.on_tuple A B (fun x => g x id) n m H xs).
Proof.
  cbv [on_tuple_cps Tuple.on_tuple]; intros H.
  rewrite to_list_cps_correct, Hg, from_list_default_cps_correct.
  rewrite (Tuple.from_list_default_eq _ _ _ (H _ (Tuple.length_to_list _))).
  reflexivity.
Qed. #[global] Hint Rewrite @on_tuple_cps_correct using (intros; autorewrite with uncps; auto): uncps.

Fixpoint update_nth_cps {A} n (g:A->A) xs {T} (f:list A->T) :=
  match n with
  | O =>
    match xs with
    | [] => f []
    | x' :: xs' => f (g x' :: xs')
    end
  | S n' =>
    match xs with
    | [] => f []
    | x' :: xs' => update_nth_cps n' g xs' (fun r => f (x' :: r))
    end
  end.
Lemma update_nth_cps_correct {A} n g: forall xs T f,
    @update_nth_cps A n g xs T f = f (update_nth n g xs).
Proof. induction n; intros; simpl; break_match; try apply IHn; reflexivity. Qed.
#[global]
Hint Rewrite @update_nth_cps_correct : uncps.

Fixpoint combine_cps {A B} (la :list A) (lb : list B)
         {T} (f:list (A*B)->T) :=
  match la with
  | nil => f nil
  | a :: tla =>
    match lb with
    | nil => f nil
    | b :: tlb => combine_cps tla tlb (fun lab => f ((a,b)::lab))
    end
  end.
Lemma combine_cps_correct {A B} la: forall lb {T} f,
    @combine_cps A B la lb T f = f (combine la lb).
Proof.
  induction la; simpl combine_cps; simpl combine; intros;
    try break_match; try apply IHla; reflexivity.
Qed.
#[global]
Hint Rewrite @combine_cps_correct: uncps.

(* differs from fold_right_cps in that the functional argument `g` is also a CPS function *)
Local Set Universe Polymorphism.
Definition fold_right_cps2_specialized_step
           (fold_right_cps2_specialized
            : forall {T A B} (g : B -> A -> (A->T)->T) (a0 : A) (l : list B) (f : A -> T), _)
           {T A B} (g : B -> A -> (A->T)->T) (a0 : A) (l : list B) (f : A -> T) :=
  match l with
  | nil => f a0
  | b :: tl => fold_right_cps2_specialized g a0 tl (fun r => g b r f)
  end.
Fixpoint fold_right_cps2_specialized {T A B} (g : B -> A -> (A->T)->T) (a0 : A) (l : list B) (f : A -> T) :=
  @fold_right_cps2_specialized_step (@fold_right_cps2_specialized) T A B g a0 l f.
Definition fold_right_cps2 {A B} (g : B -> A -> forall {T}, (A->T)->T) (a0 : A) (l : list B) {T} (f : A -> T) :=
  @fold_right_cps2_specialized T A B (fun b a => @g b a T) a0 l f.
Local Unset Universe Polymorphism.
Lemma unfold_fold_right_cps2 {A B} (g : B -> A -> forall {T}, (A->T)->T) (a0 : A) (l : list B) {T} (f : A -> T)
  : @fold_right_cps2 A B g a0 l T f
    = match l with
      | nil => f a0
      | b :: tl => fold_right_cps2 g a0 tl (fun r => g b r f)
      end.
Proof. destruct l; reflexivity. Qed.
Lemma fold_right_cps2_correct {A B} g a0 l : forall {T} f,
  (forall b a T h, @g b a T h = h (@g b a A id)) ->
  @fold_right_cps2 A B g a0 l T f = f (List.fold_right (fun b a => @g b a A id) a0 l).
Proof.
  induction l as [|?? IHl]; intros T f H; [reflexivity|].
  rewrite unfold_fold_right_cps2. simpl fold_right.
  rewrite H; erewrite IHl by eassumption.
  rewrite H; reflexivity.
Qed.
#[global]
Hint Rewrite @fold_right_cps2_correct using (intros; autorewrite with uncps; auto): uncps.

Definition fold_right_no_starter {A} (f:A->A->A) ls : option A :=
  match ls with
  | nil => None
  | cons x tl => Some (List.fold_right f x tl)
  end.
Lemma fold_right_min ls x :
  x = List.fold_right Z.min x ls
  \/ List.In (List.fold_right Z.min x ls) ls.
Proof.
  induction ls; intros; simpl in *; try tauto.
  match goal with |- context [Z.min ?x ?y] =>
                  destruct (Z.min_spec x y) as [[? Hmin]|[? Hmin]]
  end; rewrite Hmin; tauto.
Qed.
Lemma fold_right_no_starter_min ls : forall x,
    fold_right_no_starter Z.min ls = Some x ->
    List.In x ls.
Proof.
  cbv [fold_right_no_starter]; intros; destruct ls; try discriminate.
  inversion H; subst; clear H.
  destruct (fold_right_min ls z);
    simpl List.In; tauto.
Qed.
Fixpoint fold_right_cps {A B} (g:B->A->A) (a0:A) (l:list B) {T} (f:A->T) :=
  match l with
  | nil => f a0
  | cons a tl => fold_right_cps g a0 tl (fun r => f (g a r))
  end.
Lemma fold_right_cps_correct {A B} g a0 l: forall {T} f,
    @fold_right_cps A B g a0 l T f = f (List.fold_right g a0 l).
Proof. induction l as [|? l IHl]; intros; simpl; rewrite ?IHl; auto. Qed.
#[global]
Hint Rewrite @fold_right_cps_correct : uncps.

Definition fold_right_no_starter_cps {A} g ls {T} (f:option A->T) :=
  match ls with
  | nil => f None
  | cons x tl => f (Some (List.fold_right g x tl))
  end.
Lemma fold_right_no_starter_cps_correct {A} g ls {T} f :
  @fold_right_no_starter_cps A g ls T f = f (fold_right_no_starter g ls).
Proof.
  cbv [fold_right_no_starter_cps fold_right_no_starter]; break_match; reflexivity.
Qed.
#[global]
Hint Rewrite @fold_right_no_starter_cps_correct : uncps.

Import Tuple.

Module Tuple.
  Fixpoint map_cps {A B n} (g : A->B) (t : tuple A n) {T} :
    (tuple B n->T) -> T:=
  match n return tuple A n -> (tuple B n -> T) -> T with
  | O => fun _ f => f tt
  | S n' => fun t f => map_cps g (tl t) (fun r => f (append (g (hd t)) r))
  end t.
  Lemma map_cps_correct {A B n} g t: forall {T} f,
      @map_cps A B n g t T f = f (map g t).
  Proof.
    induction n; simpl map_cps; intros; try destruct t;
    [|rewrite IHn, <-map_append,<-subst_append]; reflexivity.
  Qed. #[global] Hint Rewrite @map_cps_correct : uncps.

  Fixpoint map2_cps {n A B C} (g:A->B->C) (xs:tuple A n) (ys:tuple B n) {T} :
    (tuple C n->T) -> T :=
    match n return tuple _ n -> tuple _ n -> (tuple C n -> T) -> T with
    | O => fun _ _ f => f tt
    | S n' => fun xs ys f =>
                map2_cps g (tl xs) (tl ys) (fun zs => f (append (g (hd xs) (hd ys)) zs))
    end xs ys.
  Lemma map2_cps_correct {n A B C} g xs ys : forall {T} f,
      @map2_cps n A B C g xs ys T f = f (map2 g xs ys).
  Proof.
    induction n; simpl map2_cps; intros; try destruct xs, ys;
    [|rewrite IHn, <-map2_append,<-!subst_append]; reflexivity.
  Qed. #[global] Hint Rewrite @map2_cps_correct : uncps.

  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
    := fun ts {R} => @mapi_with'_cps_specialized R T A B (fun n t a => @f n t a R) n i start ts.

  Definition mapi_with_cps {S A B n}
          (f: nat->S->A->forall {T}, (S*B->T)->T) (start:S) (ys:tuple A n) {T}
    : (S * tuple B n->T)->T :=
  match n as n0 return (tuple A n0 -> (S * tuple B n0->T)->T) with
  | O => fun ys ret => ret (start, tt)
  | S n' => fun ys ret => mapi_with'_cps 0%nat f start ys ret
  end ys.

  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).
  Proof. induction n as [|n IHn]; intros i f start xs T ret H; simpl; rewrite H, ?IHn by assumption; reflexivity. Qed.
  Lemma mapi_with_cps_correct {S A B n} f start xs T ret
  (H: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 f start xs T ret = ret (mapi_with (fun i s a => f i s a _ id) start xs).
  Proof. destruct n; simpl; rewrite ?mapi_with'_cps_correct by assumption; reflexivity. Qed.
#[global]
  Hint Rewrite @mapi_with_cps_correct @mapi_with'_cps_correct
       using (intros; autorewrite with uncps; auto): uncps.

  Section internal_mapi_with_cps2.
    (* 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'_cps2_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'_cps2_specialized (S i) (fst sb) (tl ys)
                      (fun r => ret (fst r, (snd r, snd sb))))
      end.
  End internal_mapi_with_cps2.

  Definition mapi_with'_cps2 {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
    := fun ts {R} => @mapi_with'_cps2_specialized R T A B (fun n t a => @f n t a R) n i start ts.

  Definition mapi_with_cps2 {S A B n}
          (f: nat->S->A->forall {T}, (S*B->T)->T) (start:S) (ys:tuple A n) {T}
    : (S * tuple B n->T)->T :=
  match n as n0 return (tuple A n0 -> (S * tuple B n0->T)->T) with
  | O => fun ys ret => ret (start, tt)
  | S n' => fun ys ret => mapi_with'_cps2 0%nat f start ys ret
  end ys.

  Lemma unfold_mapi_with'_cps2 {T A B n} i
        (f: nat->T->A->forall {R}, (T*B->R)->R) (start:T)
    : @mapi_with'_cps2 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'_cps2 (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'_cps2_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'_cps2 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).
  Proof. induction n as [|n IHn]; intros i f start xs T ret H; simpl; rewrite H, ?IHn by assumption; reflexivity. Qed.
  Lemma mapi_with_cps2_correct {S A B n} f start xs T ret
  (H:forall i s a R (ret:_->R), f i s a R ret = ret (f i s a _ id))
  : @mapi_with_cps2 S A B n f start xs T ret = ret (mapi_with (fun i s a => f i s a _ id) start xs).
  Proof. destruct n; simpl; rewrite ?mapi_with'_cps2_correct by assumption; reflexivity. Qed.
#[global]
  Hint Rewrite @mapi_with_cps2_correct @mapi_with'_cps2_correct
       using (intros; autorewrite with uncps; auto): uncps.

  Fixpoint left_append_cps {A n} (x:A) (xs:tuple A n) {R} :
    (tuple A (S n) -> R) -> R :=
  match
    n as n0 return (tuple A n0 -> (tuple A (S n0) -> R) -> R)
  with
  | 0%nat => fun _ f => f x
  | S n' =>
      fun xs f =>
      left_append_cps x (tl xs) (fun r => f (append (hd xs) r))
  end xs.
  Lemma left_append_cps_correct A n x xs R f :
    @left_append_cps A n x xs R f = f (left_append x xs).
  Proof.
    induction n; [reflexivity|].
    simpl left_append. simpl left_append_cps.
    rewrite IHn. reflexivity.
  Qed.

  Definition tl_cps {A n} (xs : tuple A (S n)) {R} :
    (tuple A n -> R) -> R :=
  match
    n as n0 return (tuple A (S n0) -> (tuple A n0 -> R) -> R)
  with
  | 0%nat => fun _ f => f tt
  | S n' => fun xs f => f (fst xs)
  end xs.
  Lemma tl_cps_correct A n xs R f :
    @tl_cps A n xs R f = f (tl xs).
  Proof. destruct n; reflexivity. Qed.

  Definition hd_cps {A n} (xs : tuple A (S n)) {R} :
    (A -> R) -> R :=
  match
    n as n0 return (tuple A (S n0) -> (A -> R) -> R)
  with
  | 0%nat => fun x f => f x
  | S n' => fun xs f => f (snd xs)
  end xs.
  Lemma hd_cps_correct A n xs R f :
    @hd_cps A n xs R f = f (hd xs).
  Proof. destruct n; reflexivity. Qed.

  Fixpoint left_tl_cps {A n} (xs : tuple A (S n)) {R} :
    (tuple A n -> R) -> R :=
  match
    n as n0 return (tuple A (S n0) -> (tuple A n0 -> R) -> R)
  with
  | 0%nat => fun _ f => f tt
  | S n' =>
      fun xs f =>
        tl_cps xs (fun xtl => hd_cps xs (fun xhd =>
          left_tl_cps xtl (fun r => f (append xhd r))))
  end xs.
  Lemma left_tl_cps_correct A n xs R f :
    @left_tl_cps A n xs R f = f (left_tl xs).
  Proof.
    induction n; [reflexivity|].
    simpl left_tl. simpl left_tl_cps.
    rewrite IHn. reflexivity.
  Qed.

  Fixpoint left_hd_cps {A n} (xs : tuple A (S n)) {R} :
    (A -> R) -> R :=
  match
    n as n0 return (tuple A (S n0) -> (A -> R) -> R)
  with
  | 0%nat => fun x f => f x
  | S n' =>
      fun xs f => tl_cps xs (fun xtl => left_hd_cps xtl f)
  end xs.
  Lemma left_hd_cps_correct A n xs R f :
    @left_hd_cps A n xs R f = f (left_hd xs).
  Proof.
    induction n; [reflexivity|].
    simpl left_hd. simpl left_hd_cps.
    rewrite IHn. reflexivity.
  Qed.

  Lemma In_left_hd {A} n (p : tuple A (S n))
    : In (left_hd p) (to_list _ p).
  Proof.
    simpl in *.
    induction n as [|n IHn].
    { left; reflexivity. }
    { destruct p; simpl in *; right; auto. }
  Qed.

  Lemma In_to_list_left_tl {A} n (p : tuple A (S n)) x
    : In x (to_list n (left_tl p)) -> In x (to_list (S n) p).
  Proof.
    simpl in *.
    remember (to_list n (left_tl p)) as ls eqn:H.
    intro Hin; revert Hin H.
    revert n p.
    induction ls as [|l ls IHls], n.
    { simpl; intros ? []. }
    { simpl; intros ? []. }
    { simpl; congruence. }
    { simpl; intros [? p] [H0|H1] H; simpl in *;
        setoid_rewrite to_list_append in H;
        inversion H; clear H; subst; auto. }
  Qed.

  Lemma to_list_left_append {A} {n} (p:tuple A n) x
    : (to_list _ (left_append x p)) = to_list n p ++ x :: nil.
  Proof.
    destruct n as [|n]; simpl in *; [ reflexivity | ].
    induction n as [|n IHn]; simpl in *; [ reflexivity | ].
    destruct p; simpl in *; rewrite IHn; simpl; reflexivity.
  Qed.
End Tuple.
#[global]
Hint Rewrite @Tuple.map_cps_correct @Tuple.left_append_cps_correct @Tuple.left_tl_cps_correct @Tuple.left_hd_cps_correct @Tuple.tl_cps_correct @Tuple.hd_cps_correct : uncps.
#[global]
Hint Rewrite @Tuple.mapi_with_cps_correct @Tuple.mapi_with'_cps_correct @Tuple.mapi_with_cps2_correct @Tuple.mapi_with'_cps2_correct
     using (intros; autorewrite with uncps; auto): uncps.
