Library Containers.Decorate_M

Require Export Containers.Container.
Require Import Containers.Final_Coalgebra.

Section Decorate_M.

Variable funext : Funext.

CoInductive M (A : Type) (B : A Type) :=
sup (a : A) (f : B a M A B).
Arguments sup {_ _} _ _.

Definition label {A B} (m : M A B) : A :=
match m with sup a fa end.

Definition arg {A B} (m : M A B) : B (label m) M A B :=
match m with sup a ff end.

Definition out {A B} (m : M A B) : a, B a M A B :=
(label m; arg m).

Instance is_equiv_out {A B} : IsEquiv (@out A B).
Proof.
  serapply BuildIsEquiv.
  - intros [a f].
    exact (sup a f).
  - intros [a f].
    reflexivity.
  - intros [a f].
    reflexivity.
  - intros [a f].
    reflexivity.
Defined.

Definition m_coalg A B := Build_Coalgebra (| A ||> B |) (M A B) out.

CoFixpoint corec {A B X}
    (x_label : X A) (x_arg : x, B (x_label x) X) (x : X) : M A B :=
sup (x_label x) (corec x_label x_arg o x_arg x).

Definition m_path {A B} (m1 m2 : M A B) :=
p : label m1 = label m2, b, arg m1 b = arg m2 (p # b).
Infix "=M" := m_path (at level 50).

Definition equiv_path_m' {A B} (m1 m2 : M A B) :
  ( p : label m1 = label m2,
     transport (fun aB a M A B ) p (arg m1) = arg m2) <~>
  (m1 = m2).
Proof.
  serapply equiv_adjointify.
  - destruct m1, m2. simpl in ×.
    intros [p q]. destruct p, q.
    reflexivity.
  - intros p.
     (ap label p).
    exact ((transport_compose (fun aB a M A B) label p (arg m1))^ @
           apD arg p).
  - intros [].
    destruct m1, m2.
    reflexivity.
  - destruct m1, m2. simpl in ×.
    intros [p q]. destruct p, q.
    reflexivity.
Defined.

Definition equiv_path_m {A B} (m1 m2 : M A B) :
  (m1 =M m2) <~>
  (m1 = m2).
Proof.
  refine (transitivity _ (equiv_path_m' m1 m2)).
  apply equiv_functor_sigma_id; intros p.
  refine (transitivity _ (equiv_moveR_transport_p _ _ _ _)).
  refine (transitivity _ (equiv_inverse (equiv_ap10 _ _))).
  apply equiv_functor_forall_id; intros b.
  apply (equiv_concat_r).
  rewrite transport_arrow_toconst.
  rewrite inv_V.
  reflexivity.
Defined.

Definition m_path_label {A B} {m1 m2 : M A B} (p : m1 = m2) :
  label m1 = label m2 :=
((equiv_path_m m1 m2)^-1 p).1.

Definition m_path_arg {A B} {m1 m2 : M A B} (p : m1 = m2) (b : B (label m1)) :
  arg m1 b = arg m2 (m_path_label p # b) :=
((equiv_path_m m1 m2)^-1 p).2 b.

Axiom is_final_m : A B, is_final (m_coalg A B).

Axiom eta :
   {A B X}
    (f : X M A B) (g : X M A B)
    (x_arg : x, B (label (f x)) X)
    (p : x, label (f x) = label (g x))
    (q1 : x b, arg (f x) b = f (x_arg x b))
    (q2 : x b, arg (g x) (p x # b) = g (x_arg x b)),
   x, f x = g x.

Axiom destruct_eta :
   {A B X}
    (f : X M A B) (g : X M A B)
    (x_arg : x, B (label (f x)) X)
    (p : x, label (f x) = label (g x))
    (q1 : x b, arg (f x) b = f (x_arg x b))
    (q2 : x b, arg (g x) (p x # b) = g (x_arg x b))
    (x : X),
  (m_path_label (eta f g x_arg p q1 q2 x);
   m_path_arg (eta f g x_arg p q1 q2 x)) =
  existT (fun s b, arg (f x) b = arg (g x) (s # b))
         (p x)
         (fun bq1 x b @ eta f g x_arg p q1 q2 (x_arg x b) @ (q2 x b)^).

Inductive Addr {A B} (m : M A B) : A Type :=
| root_addr : Addr m (label m)
| subtree_addr a b : Addr (arg m b) a Addr m a.
Arguments root_addr {_ _ _}.
Arguments subtree_addr {_ _ _ _} _ _.

Fixpoint transport_addr {A B} {m1 m2 : M A B} (p : m1 = m2)
    {a} (addr : Addr m1 a) {struct addr} :
  Addr m2 a :=
match addr with
  | root_addrtransport (fun aAddr m2 a) (m_path_label p)^ root_addr
  | subtree_addr _ b addr'
    subtree_addr (transport B (m_path_label p) b)
                 (transport_addr (m_path_arg p b) addr')
end.

Lemma transport_addr_correct {A B} {m1 m2 : M A B} (p : m1 = m2)
    {a} (addr : Addr m1 a) :
  transport (fun m'Addr m' a) p addr =
  transport_addr p addr.
Proof.
  destruct p.
  clear.
  induction addr.
  - reflexivity.
  - exact (ap (subtree_addr b) IHaddr).
Qed.

Variables A1 : Type.
Variable A2 B : A1 Type.

Definition dec_step_label (mf : {m : M A1 B & a, Addr m a A2 a}) :=
existT A2 (label mf.1) (mf.2 _ root_addr).

Definition dec_step_arg (mf : {m : M A1 B & a, Addr m a A2 a}) :=
fun bexistT (fun m a, Addr m a A2 a)
              (arg mf.1 b)
              (fun amf.2 a o subtree_addr b).

Definition dec :
  {m : M A1 B & a, Addr m a A2 a}
  M {a : A1 & A2 a} (B o pr1) :=
corec dec_step_label dec_step_arg.

Definition undec1_step_label (m : M {a : A1 & A2 a} (B o pr1)) :=
pr1 (label m).

Definition undec1_step_arg (m : M {a : A1 & A2 a} (B o pr1)) :=
arg m.

Definition undec1 : M {a : A1 & A2 a} (B o pr1) M A1 B :=
corec undec1_step_label undec1_step_arg.

Fixpoint undec2 (m : M {a : A1 & A2 a} (B o pr1))
    a (addr : Addr (undec1 m) a) {struct addr} :
  A2 a.
Proof.
  destruct addr.
  - exact (pr2 (label m)).
  - exact (undec2 (arg m b) _ addr).
Defined.

Definition undec m : {m : _ & a, Addr m a A2 a} :=
(undec1 m; undec2 m).

Definition dec_undec m :
  dec (undec m) = m :=
eta (dec o undec) idmap undec1_step_arg
  (fun _idpath) (fun _ _idpath) (fun _ _idpath) m.

Definition undec1_dec m f :
  m = undec1 (dec (m; f)) :=
eta pr1 (undec1 o dec) dec_step_arg
  (fun _idpath) (fun _ _idpath) (fun _ _idpath) (m; f).

Fixpoint transport_addr_undec1_dec m f {a} (addr : Addr m a) :
  Addr (undec1 (dec (m; f))) a :=
match addr with
  | root_addrroot_addr
  | subtree_addr _ b addr'
    @subtree_addr _ _ (undec1 (dec (m; f))) _ b
                  (transport_addr_undec1_dec (arg m b)
                                             (fun af a o subtree_addr b) addr')
end.

Lemma transport_addr_undec1_dec_correct m f {a} (addr : Addr m a) :
transport (fun mAddr m a) (undec1_dec m f) addr =
transport_addr_undec1_dec m f addr.
Proof.
  apply (transitivity (transport_addr_correct (undec1_dec m f) addr)).
  induction addr; simpl.
  - unfold undec1_dec.
    apply (ap
        (fun ptransport (fun aAddr (undec1 (dec (m; f))) a) p^ root_addr)
        (destruct_eta pr1 (undec1 o dec) dec_step_arg
                      (fun _idpath) (fun _ _idpath) (fun _ _idpath)
                      (m; f))..1).
  - rewrite <- IHaddr.
    transitivity (
        (fun p : m_path _ _
           subtree_addr (transport B p.1 b)
                        (transport_addr (p.2 b) addr))
          (m_path_label (undec1_dec m f); m_path_arg (undec1_dec m f))). {
      reflexivity.
    }
    transitivity (
        (fun p : m_path m (undec1 (dec (m; f))) ⇒
           subtree_addr (transport B p.1 b)
                        (transport_addr (p.2 b) addr))
          (existT (fun p b, arg m b = arg (undec1 (dec (m; f))) (p # b))
                  idpath
                  (fun bundec1_dec (arg m b) (fun af a o subtree_addr b)))). {
      f_ap.
      unfold undec1_dec.
      apply (transitivity
               (destruct_eta pr1 (undec1 o dec) dec_step_arg
                             (fun _idpath) (fun _ _idpath) (fun _ _idpath)
                             (m; f))).
      f_ap.
      apply path_forall; intros b'.
      autorewrite with paths.
      reflexivity.
    }
    reflexivity.
Qed.

Lemma undec2_dec m f :
  transport (fun m' a, Addr m' a A2 a) (undec1_dec m f) f =
  (undec2 (dec (m; f))).
Proof.
  apply path_forall; intros a.
  rewrite transport_forall_constant.
  apply (dpath_arrow (fun m'Addr m' a) (fun _A2 a));
    intros addr.
  rewrite transport_const.
  rewrite transport_addr_undec1_dec_correct.
  induction addr.
  - simpl.
    reflexivity.
  - simpl.
    exact (IHaddr (fun af a o subtree_addr b)).
Qed.

Lemma undec_dec m f :
  (m; f) = undec (dec (m; f)).
Proof.
  serapply path_sigma'.
  - apply undec1_dec.
  - apply undec2_dec.
Qed.

Lemma equiv_decorate_M :
  M {a : A1 & A2 a} (B o pr1) <~> {m : M A1 B & a, Addr m a A2 a}.
Proof.
  serapply equiv_adjointify.
  - exact undec.
  - exact dec.
  - intros [m f].
    symmetry.
    apply undec_dec.
  - intros m.
    apply dec_undec.
Qed.

End Decorate_M.