A formulation for W-types which supports the correct recursion principle on nat

W.v

(** Some imports *)
Require Import Coq.Vectors.Vector.
Require Import Coq.Vectors.Fin.
Require Import Coq.Logic.FunctionalExtensionality. (* only used for qinv equalities, not recursor *)
Set Primitive Projections.
Set Universe Polymorphism.
(** In [W A B], [A] labels the constructors and their non-recursive
    arguments, while [B] describes the recursive arguments for each
    constructor. *)
Inductive W (A : Type) (B : A -> Type) := sup (wnr : A) (wr : B wnr -> W A B).
Arguments sup {_ _} _ _.
(** Unfortunately, constructing the right eliminator often requires
    function extensionality.  For inductive types which don't use
    functions in their constructors, we can construct a variant of [W]
    types without this issue.  We construct a variant that splits the
    non-recursive and recursive arguments apart, so we don't need
    functions out of empty types, and we split apart finite recursion
    from labeled recursion so we don't need functions where we
    wouldn't for normal inductive types *)
(* inductive version of [match] for positivity checker *)
Inductive match_option {A} (S : A -> Type) (N : Type) : option A -> Type :=
| match_Some {x} : S x -> match_option S N (Some x)
| match_None     : N   -> match_option S N None.
Arguments match_Some {_ _ _ _} _.
Arguments match_None {_ _ _} _.
Inductive W' : forall (A_nr : option Type)
                      (A_r : option Type)
                      (B : match A_r return Type with
                           | None => unit
                           | Some A_r => A_r -> nat (* finite recursion count *) * option Type (* non-finite recursion labels *)
                           end),
    Type :=
| sup_nr {A_nr A_r B} (W:=W' (Some A_nr) A_r B) (_ : A_nr) : W
| sup_r {A_nr A_r B} (W:=W' A_nr (Some A_r) B)
        (r : A_r)
        (_ : Vector.t W (fst (B r)))
        (_ : match_option (fun T => T -> W) unit (snd (B r))) : W.
Arguments Vector.nil {_}.
Arguments Vector.cons {A} _ {n} _.
Section Vector_Forall.
  Context {A} (P : A -> Type).

  Fixpoint Vector_Forall {n} : Vector.t A n -> Type
    := match n return Vector.t A n -> Type with
       | 0 => fun _ => unit
       | S n => fun v => P (Vector.hd v) * Vector_Forall (Vector.tl v)
       end%type.

  Section apply.
    Context (p : forall a, P a).

    Fixpoint Vector_Forall_forall {n} (v : Vector.t A n) : Vector_Forall v
      := match v return Vector_Forall v with
         | nil => tt
         | cons h x => (p _, @Vector_Forall_forall _ x)
         end.
  End apply.
End Vector_Forall.
Arguments Vector_Forall_forall {A P} _ {n} v.
Section Vector_Forall_nth.
  Context {A} {P : A -> Type}.
  Fixpoint Vector_Forall_nth {n v} : forall p, @Vector_Forall A P n v -> P (Vector.nth v p)
    := match v return forall p, @Vector_Forall A P _ v -> P (Vector.nth v p) with
       | nil => ltac:(now inversion 1)
       | cons v vs
         => fun p H
            => match p in Fin.t n
                     return forall vs : Vector.t A (pred n),
                   (forall p, P (Vector.nth vs p))
                   -> match n return Fin.t n -> Vector.t A (pred n) -> Type with
                      | 0 => fun _ _ => unit
                      | S n => fun p vs => P (Vector.nth (cons v vs) p)
                      end p vs
               with
               | F1 => fun _ _ => fst H
               | FS p => fun _ H => H _
               end vs (fun p => @Vector_Forall_nth _ vs p (snd H))
       end.

  Lemma Vector_Forall_nth_forall {n v p H}
    : @Vector_Forall_nth n v p (Vector_Forall_forall H v) = H _.
  Proof.
    induction v.
    { inversion p. }
    { revert IHv.
      revert v.
      revert p.
      revert n.
      lazymatch goal with
      | [ |- forall n (p : Fin.t (S n)), @?P n p ]
        => refine (fun n p => match p as p' in Fin.t n'
                                    return match n' return Fin.t n' -> Type with
                                           | 0 => fun _ => unit
                                           | S n' => fun p' => P n' p'
                                           end p'
                              with
                              | Fin.F1 => _
                              | Fin.FS _ => _
                              end)
      end.
      { reflexivity. }
      { cbn; auto. } }
  Qed.
End Vector_Forall_nth.
Section Vector_make.
  Context {A : Type}.
  Fixpoint Vector_make {n} : (Fin.t n -> A) -> Vector.t A n
    := match n return (Fin.t n -> A) -> Vector.t A n with
       | 0 => fun _ => nil
       | S n => fun f => cons (f F1) (@Vector_make n (fun fv => f (FS fv)))
       end.
  Lemma nth_Vector_make {n H} p : Vector.nth (@Vector_make n H) p = H p.
  Proof.
    induction p; cbn; [ reflexivity | apply IHp ].
  Qed.
End Vector_make.

Fixpoint W'_rect'
         (BT := fun A_r => match A_r return Type with
                           | None => unit
                           | Some A_r => A_r -> nat (* finite recursion count *) * option Type (* non-finite recursion labels *)
                           end)
         {A_nr : option Type}
         {A_r : option Type}
         {B : BT A_r}
         (W:=W' A_nr A_r B)
         (P : W -> Type)
         (Pnr : match A_nr return (W' A_nr A_r B -> Type) -> Type with
                | Some A_nr => fun P => forall v : A_nr, P (sup_nr v)
                | None => fun _ => unit
                end P)
         (Pr : match A_r return forall B : BT A_r, (W' A_nr A_r B -> Type) -> Type with
               | Some A_r
                 => fun B P
                    => let W:=W' A_nr (Some A_r) B in
                       forall (r : A_r)
                              (rf : Vector.t W (fst (B r)))
                              (rfP : Vector_Forall P rf)
                              (rnf : match_option (fun T => T -> W) unit (snd (B r)))
                              (rnfP : match rnf with
                                      | match_None _ => unit
                                      | match_Some rnf => forall t, P (rnf t)
                                      end),
                         P (sup_r r rf rnf)
               | None => fun _ _ => unit
               end B P)
         (v : W)
  : P v.
Proof.
  destruct v as [|A_nr A_r B W0 r rf rnf].
  { apply Pnr. }
  { specialize (@W'_rect' A_nr (Some A_r) B P Pnr Pr).
    apply Pr.
    { apply Vector_Forall_forall; intros; apply W'_rect'. }
    { destruct rnf; try exact tt; intros; apply W'_rect'. } }
Defined.

Record qinv A B := { fwd :> A -> B ; bak : B -> A ; fwd_bak : forall x, fwd (bak x) = x ; bak_fwd : forall x, bak (fwd x) = x }.
Local Notation iffT A B := (prod (A -> B) (B -> A)).
Import EqNotations.

Definition W_of_W' {A_nr A_r B}
  : qinv (W' A_nr A_r B)
         (W (match A_nr with Some T => T | None => Empty_set end
             + match A_r with Some T => T | None => Empty_set end)
            (fun a => match a with
                      | inl a => Empty_set
                      | inr a
                        => match A_r
                                 return match A_r with Some T => T | None => Empty_set end
                                        -> match A_r return Type with
                                           | None => unit
                                           | Some A_r => A_r -> nat (* finite recursion count *) * option Type (* non-finite recursion labels *)
                                           end
                                        -> Type
                           with
                           | Some T => fun a B => Fin.t (fst (B a)) + match snd (B a) with Some T => T | None => Empty_set end
                           | None => fun _ _ => Empty_set
                           end%type a B
                      end)).
Proof.
  revert A_nr A_r B.
  simple refine (let fwd := _ in
                 let bak := _ in
                 fun A_nr A_r B => @Build_qinv _ _ (fwd A_nr A_r B) (bak A_nr A_r B) _ _);
    [ intros A_nr A_r B | intros A_nr A_r B | | ].
  { refine (W'_rect' _ _ _).
    { destruct A_nr; try exact tt.
      refine (fun v => sup (inl v) (fun x : Empty_set => match x with end)). }
    { destruct A_r as [A_r|]; cbv beta iota zeta; try exact tt.
      intros r rf Hrf rnf Hnf.
      apply (sup (inr r)).
      intros [n|l].
      { refine (Vector_Forall_nth n Hrf). }
      { destruct rnf; [ | case l ].
        refine (Hnf l). } } }
  { clear fwd.
    fix W_of_W'_rev 1.
    refine (fun w => match w with
                     | sup wnr wr => _
                     end); clear w; destruct wnr as [wnr|wnr].
    { clear W_of_W'_rev.
      destruct A_nr as [A_nr|]; [ apply sup_nr; exact wnr | exfalso; destruct wnr ]. }
    { specialize (fun x => W_of_W'_rev (wr x)); clear wr.
      destruct A_r as [A_r|]; [ apply (sup_r wnr) | exfalso; destruct wnr ].
      { refine (Vector_make (fun n => W_of_W'_rev (inl n))). }
      { destruct (snd (B wnr)); constructor; try exact tt.
        refine (fun l => W_of_W'_rev (inr l)). } } }
  { fix IH 1; intros [[wnr|wnr] wr].
    { clear IH.
      destruct A_nr as [A_nr|]; cbn; [ | destruct wnr ].
      apply f_equal; extensionality H; exfalso; case H. }
    { specialize (fun x => IH (wr x)).
      destruct A_r as [A_r|]; [ | exfalso; destruct wnr ].
      cbn; apply f_equal.
      extensionality l.
      destruct l as [n|l]; [ | destruct (snd (B wnr)); [ apply IH | case l ] ].
      etransitivity; [ | eapply IH ]; clear IH.
      cbv [fwd].
      rewrite !Vector_Forall_nth_forall, !nth_Vector_make.
      reflexivity. } }
  { refine (W'_rect' _ _ _).
    { destruct A_nr; try exact tt; reflexivity. }
    { destruct A_r as [A_r|]; cbv beta iota zeta; try exact tt.
      intros r rf Hrf rnf Hrnf.
      cbn.
      apply f_equal2.
      { apply eq_nth_iff.
        intros; subst.
        change (@VectorDef.nth) with (@Vector.nth).
        rewrite !nth_Vector_make, !Vector_Forall_nth_forall.
        pose proof (fun p => Vector_Forall_nth p Hrf) as Hrf'.
        cbv beta in *.
        apply Hrf'. }
      { destruct rnf; repeat match goal with u : unit |- _ => destruct u end; [ | reflexivity ].
        apply f_equal.
        extensionality l.
        cbn.
        etransitivity; [ | apply Hrnf ].
        reflexivity. } } }
Defined.

Definition W'_of_W {A B} : qinv (W A B) (W' None (Some A) (fun a => (0, Some (B a)))).
Proof.
  revert A B.
  simple refine (let fwd := _ in
                 let bak := _ in
                 fun A B => @Build_qinv _ _ (fwd A B) (bak A B) _ _);
    [ intros A B | intros A B | | ].
  { fix W'_of_W 1.
    refine (fun w => match w with
                     | sup wnr wr => @sup_r None A (fun a => (0, Some (B a))) wnr nil (match_Some (fun b => @W'_of_W (wr _)))
                     end).
    refine b. }
  { refine (W'_rect' _ _ _); cbn [fst snd Vector_Forall]; try exact tt.
    intros r _ _ rnf Hrnf.
    refine (sup r _).
    revert Hrnf.
    refine match rnf in match_option _ _ v return match v with Some T => _ -> T -> _ | None => unit end with
           | match_Some _ => _
           | match_None _ => tt
           end.
    clear; auto. }
  { refine (W'_rect' _ _ _); cbn [fst snd fwd bak]; try exact tt.
    intros r ? ? rnf Hrnf.
    assert (H' : exists v, rnf = match_Some v).
    { clear.
      refine match rnf with match_Some _ => _ | match_None _ => tt end.
      eexists; reflexivity. }
    { destruct H' as [v ?]; subst.
      refine match rf with nil => _ | cons _ _ => tt end.
      cbn.
      do 2 apply f_equal.
      extensionality x.
      apply Hrnf. } }
  { fix IH 1; intros [x]; cbn in *.
    apply f_equal; extensionality b; apply IH. }
Defined.

Definition nat := W' (Some unit) (Some unit) (fun _ => (1, None)).
Definition O : nat := sup_nr tt.
Definition S (x : nat) : nat := sup_r tt (cons x nil) (match_None tt).
Definition nat_rect (P : nat -> Type) (fO : P O) (fS : forall n, P n -> P (S n)) (n : nat) : P n.
Proof.
  refine (W'_rect' P (fun 'tt => fO) (fun 'tt vn '(x, tt) m _ => _ (fS (Vector.hd vn) x)) n).
  clear x; cbv [S]; cbn in *.
  repeat first [ exact idProp
               | refine match vn with
                        | Vector.nil => _
                        | Vector.cons v vn => _
                        end; [ clear vn | ]
               | refine match m with
                        | match_Some _ => tt
                        | match_None tt => _
                        end; clear dependent m
               | progress clear ].
  exact id.
Defined.
Check fun P fO fS => eq_refl : nat_rect P fO fS O = fO.
Check fun P fO fS n => eq_refl : nat_rect P fO fS (S n) = fS n (nat_rect P fO fS n).

Thanks for that code