arrow-monad-comonad.v

Set Implicit Arguments.
Unset Strict Implicit.
Generalizable All Variables.

Local Close Scope type_scope.

Definition id {A}: A -> A := fun x => x.
Definition id_ (A: Type) := id (A:=A).
Hint Unfold id id_.


Structure Setoid :=
  { carrier:> Type;
    equal: carrier -> carrier -> Prop;
    equal_refl x: equal x x;
    equal_sym x y: equal x y -> equal y x;
    equal_trans x y z: equal x y -> equal y z -> equal x z
  }.
Arguments equal {setoid}(x y): rename.
Notation "f == g" := (equal f g) (at level 90, no associativity).
Notation makeSetoid eq := (@Build_Setoid _ eq _ _ _).



Structure Map (X Y: Setoid) :=
  { fbody:> X -> Y;
    eqSubst: forall x y, x == y -> fbody x == fbody y }.
Notation makeMap f := (@Build_Map _ _ f _).
Notation "[ x ':->' p ]" := (makeMap (fun x => p)).

Definition eqMap {X Y: Setoid}(f g: Map X  Y): Prop :=
  forall x, fbody f x == fbody g x.


Program Definition mapSetoid (X Y: Setoid): Setoid := makeSetoid (@eqMap X Y).
Next Obligation.
  intros X Y f x; apply equal_refl.
Qed.
Next Obligation.
  intros X Y f g Heq x; apply equal_sym, Heq; auto.
Qed.
Next Obligation.
  intros X Y f g h Heqfg Heqgh x; apply equal_trans with (g x); auto.
Qed.

Notation "X --> Y" := (mapSetoid X Y) (at level 60, right associativity).


Lemma eqMapEq {X Y: Setoid}(f g: X --> Y) x:
  f == g -> f x == g x.
Proof.
  simpl; intros; auto. 
Qed.


Program Definition compMap `(f: X --> Y)`(g: Y --> Z): X --> Z := [ x :-> g (f x)].
Next Obligation.
  intros; do 2 apply eqSubst; assumption.
Qed.

Notation "g \o f" := (compMap f g) (at level 60, right associativity).
Notation "f \; g" := (g \o f) (at level 55, left associativity).

Lemma compA {X Y Z W: Setoid}(f: X --> Y)(g: Y --> Z)(h: Z --> W):
  f \; (g \; h) == (f \; g) \; h.
Proof.
  simpl; intro x; simpl; apply equal_refl.  
Qed.


Program Definition idMap {X: Setoid}: X --> X := [ x :-> x].
Next Obligation.
  intros; auto.
Qed.

Definition idMap_(X: Setoid) := (@idMap X).


Lemma compMapEq {X Y Z: Setoid}(f f': X --> Y)(g g': Y --> Z):
  f == f' -> g == g' -> f \; g == f' \; g'.
Proof.
  intros.
  simpl; intro; simpl.
  apply equal_trans with (g (f' x)).
  - apply eqSubst, H.
  - apply eqMapEq, H0.
Qed.


Program Definition prodSetoid (X Y: Setoid): Setoid :=
  makeSetoid (fun (p q: prod X Y) => and (fst p == fst q) (snd p == snd q)).
Next Obligation.
  intros; split; apply equal_refl.
Qed.
Next Obligation.
  intros; split; destruct H; apply equal_sym; auto.
Qed.
Next Obligation.
  intros X Y x y z H1 H2; split; destruct H1, H2; eapply equal_trans; eassumption.
Qed.
Notation "X * Y" := (prodSetoid X Y).


Lemma eqFst {X Y: Setoid}(p q: prodSetoid X Y): p == q -> fst p == fst q.
Proof.
  simpl; intros [Hf Hs]; auto.
Qed.

Program Definition fstMap X Y: X * Y --> X := makeMap (@fst X Y).
Next Obligation.
  intros; apply eqFst; auto.
Qed.


Lemma eqSnd {X Y: Setoid}(p q: prodSetoid X Y): p == q -> snd p == snd q.
Proof.
  simpl; intros [Hf Hs]; auto.
Qed.

Program Definition sndMap X Y: X * Y --> Y := makeMap (@snd X Y).
Next Obligation.
  intros; apply eqSnd; auto.
Qed.

Hint Resolve eqFst eqSnd.


Program Definition xMap `(f: X --> Y)`(g: Z --> W): X * Z --> Y * W :=
  [ p :-> (f (fst p), g (snd p))].
Next Obligation.
  intros; simpl; split; apply eqSubst; auto.
Qed.  

Notation "f \x g" := (xMap f g) (at level 55, left associativity).


Program Definition Pair {X Y: Setoid}: X --> Y --> X * Y :=
  [x :-> [y :-> (x, y)]].
Next Obligation.
  intros; simpl; split; auto.
  apply equal_refl.
Qed.
Next Obligation.
  intros; simpl; intro; simpl; split; auto; apply equal_refl.
Qed.

Notation "( x , y )" := (Pair x y).

Program Definition assoc {X Y Z: Setoid}: (X * Y) * Z --> X * (Y * Z) :=
  [ p :-> (fst (fst p), (snd (fst p), snd p))].
Next Obligation.
  simpl; tauto.
Qed.
Hint Unfold assoc.



Reserved Notation "f >>> g" (at level 60, right associativity).
Class Arrow (a: Setoid -> Setoid -> Setoid) :=
  { pure {X Y: Setoid}: (X --> Y) --> a X Y;
    conn {X Y Z: Setoid}: a X Y --> a Y Z --> a X Z where "f >>> g" := (conn f g);
    frst {X Y Z: Setoid}: a X Y --> a (X * Z) (Y * Z);

    conn_assoc:
      forall {X Y Z W: Setoid}(f: a X Y)(g: a Y Z)(h: a Z W),
        (f >>> g) >>> h == f >>> (g >>> h);

    arr_comp:
      forall {X Y Z: Setoid}(f: X --> Y)(g: Y --> Z),
        pure (f \; g) == pure f >>> pure g;

    arr_id_l:
      forall {X Y: Setoid}(f: a X Y),
        pure idMap >>> f == f;
    
    arr_id_r:
      forall {X Y: Setoid}(f: a X Y),
        f >>> pure idMap == f;
    
    frst_proj:
      forall {X Y Z: Setoid}(f: a X Y),
        frst f >>> pure (fstMap Y Z) == pure (fstMap X Z) >>> f;

    frst_id_times:
      forall {X Y Z: Setoid}(f: a X Y)(p: Z --> Z),
        frst f >>> pure (idMap \x p) == pure (idMap \x p) >>> frst f;

    frst_twice:
      forall {X Y Z W: Setoid}(f: a X Y),
        frst (Z:=W) (frst (Z:=Z) f) >>> pure assoc == pure assoc >>>  (frst f);

    frst_arr_times:
      forall {X Y Z: Setoid}(p: X --> Y),
        frst (pure p) == pure (p \x (idMap_ Z));
    
    frst_conn:
      forall {X Y Z W: Setoid}(f: a X Y)(g: a Y Z),
        frst (Z:=W) (f >>> g) == frst f >>> frst g
  }.



Reserved Notation "mx >>= f" (at level 55, left associativity).
Class Monad (m: Setoid -> Setoid) :=
  { unit {X: Setoid}: X --> (m X);
    bind {X Y: Setoid}: (X --> m Y) --> m X --> m Y where "mx >>= f" := (bind f mx);

    bind_unit:
      forall {X Y: Setoid}(x: X)(f: X --> m Y),
        unit x >>= f == f x;
    unit_bind:
      forall {X: Setoid}(mx: m X),
        mx >>= unit == mx;
    bind_assoc:
      forall {X Y Z: Setoid}(f: X --> m Y)(g: Y -->  m Z)(mx: m X),
        mx >>= f >>= g == mx >>= (bind g \o f) }.
Notation "mx >>= f" := (bind f mx) (at level 55, left associativity).


Program Definition flip {X Y Z: Setoid}: (X --> Y --> Z) --> (Y --> X --> Z) :=
  [ f :-> [y :-> [x :-> f x y ]]].
Next Obligation.
  simpl; intros X Y Z f y x x' Heq; simpl.
  apply eqMapEq, eqSubst; auto.
Qed.  
Next Obligation.
  simpl; intros; intro; simpl.
  apply eqSubst; auto.
Qed.
Next Obligation.
  simpl; intros; intro; simpl; intro; simpl.
  apply eqMapEq, H.
Qed.


Program Definition strength `(mon: Monad m){X Y: Setoid}: m X --> Y --> m (X * Y) :=
  flip [ y :-> bind (unit \o (flip Pair y))].
Next Obligation.
  intros.
  apply eqSubst.
  apply compMapEq.
  - apply eqSubst, H.
  - apply equal_refl.
Qed.  

Reserved Notation "wx =>> f" (at level 55, left associativity).
Class Comonad (w: Setoid -> Setoid) :=
  { counit {X: Setoid}: w X --> X;
    cobind {X Y: Setoid}: (w X --> Y) --> w X --> w Y where "wx =>> f" := (cobind f wx);

    cobind_counit:
      forall {X: Setoid}(wx: w X),
        cobind counit wx == wx;
    counit_cobind:
      forall {X Y: Setoid}(f: w X --> Y)(wx: w X),
        counit (wx =>> f) = f wx;
    cobind_assoc:
      forall {X Y Z: Setoid}(f: w X --> Y)(g: w Y -->  Z)(wx: w X),
        wx =>> f =>> g == wx =>> (g \o cobind f) }.
Notation "wx =>> f" := (cobind f wx) (at level 55, left associativity).



Definition Kleisli (m: Setoid -> Setoid)(X Y: Setoid): Setoid := mapSetoid X (m Y).

Definition CoKleisli (w: Setoid -> Setoid)(X Y: Setoid): Setoid := mapSetoid (w X) Y.

Class Dist w m :=
  { dist_comonad: Comonad w;
    dist_monad: Monad m;
    dist {X: Setoid}: w (m X) --> m (w X)
  (* distributive laws *)
  }.

Definition BiKleisli w m (X Y: Setoid) := mapSetoid (w X) (m Y).



Program Definition mmap `(_: Monad m)`(f: X --> Y): m X --> m Y := makeMap (bind (unit \o f)).
Next Obligation.
  intros; apply eqSubst; auto.
Qed.

Program Definition uncurry_b {X Y Z: Setoid}: (X --> Y --> Z) -> (X * Y --> Z) :=
  fun f => [ p :-> f (fst p) (snd p)].
Next Obligation.
  simpl; intros.
  apply equal_trans with (f (fst x) (snd y)).
  - apply eqSubst, H.
  - apply eqMapEq, eqSubst, H.
Qed.

Program Definition uncurry {X Y Z: Setoid}: (X --> Y --> Z) --> (X * Y --> Z) :=
  [ f :-> uncurry_b f].
Next Obligation.
  intros; simpl; intro; simpl.
  apply eqMapEq, eqMapEq, H.
Qed.

Program Definition conn_b {X Y Z}`(mon: Monad m): X --> m Y -> (Y --> m Z) --> X --> m Z :=
  fun f => [g :-> bind g \o f ].
Next Obligation.
  intros.
  apply compMapEq.
  - apply equal_refl.
  - apply eqSubst, H.
Qed.


(* Kleisli arrow is Arrow *)
Program Instance kleisliArrow `(mon: Monad m): Arrow (Kleisli m) :=
  { pure X Y := [ f :-> unit \o f];
    conn X Y Z := [f :-> conn_b mon f];
    frst X Y Z := [f :-> uncurry (strength mon) \o (f \x idMap)] }.
Next Obligation.
  simpl; intros m mon X Y f g Heq x; simpl.
  apply eqSubst, Heq.
Qed.
Next Obligation.
  simpl; intros m mon X Y Z f g Heq x y; simpl.
  apply eqSubst, Heq.
Qed.
Next Obligation.
  simpl; intros; intro; simpl.
  apply eqSubst, eqMapEq, H.
Qed.
Next Obligation.
  intros m mon X Y Z W f g h; simpl; intro x; simpl.
  apply bind_assoc.
Qed.
Next Obligation.
  intros m mon X Y Z f g; simpl; intro x; simpl.
  apply equal_trans with ((g \; unit) (f x)).
  - simpl; apply equal_refl.
  - apply equal_sym, bind_unit.
Qed.
Next Obligation.
  intros m mon X Y f; simpl; intro x; simpl.
  apply bind_unit.
Qed.
Next Obligation.
  intros m mon X Y f; simpl; intro x; simpl.
  apply equal_trans with (f x >>= unit).
  - apply eqMapEq, eqSubst; simpl; intro; simpl; apply equal_refl.
  - apply unit_bind.
Qed.
Next Obligation.
  intros m mon X Y Z f; simpl; intros p; simpl.
  apply equal_trans with (f (fst p)).
  - eapply equal_trans; [apply bind_assoc |].
    eapply equal_trans; [apply eqMapEq |].
    + simpl; intro my.
      eapply equal_trans; [| apply bind_assoc ].
      apply equal_sym.
      eapply equal_trans; [ apply bind_assoc |].
      apply eqMapEq, eqSubst; simpl; intro y; simpl.
      apply equal_sym.
      eapply equal_trans; [ apply bind_unit |]; simpl.
      simpl.
      apply equal_trans with (unit y >>= unit); [| apply equal_refl].
      apply equal_sym.
      apply bind_unit.
    + apply equal_sym.
      eapply equal_trans; [| apply bind_assoc].
      apply equal_sym.
      eapply equal_trans; [apply unit_bind |].
      eapply equal_trans; [apply unit_bind |].
      apply equal_refl.
  - apply equal_sym, bind_unit.
Qed.
Next Obligation.
  intros m mon X Y Z f k; simpl; intro p; simpl.
  apply equal_sym.
  eapply equal_trans; [ apply bind_unit |]; simpl.
  apply equal_sym.
  eapply equal_trans; [ apply bind_assoc |]; simpl.
  apply eqMapEq, eqSubst.
  apply equal_sym.
  eapply equal_trans; [| apply compA ].
  apply equal_sym.
  eapply equal_trans; [ apply compMapEq |].
  - apply equal_refl.
  - simpl; intro q; simpl; apply bind_unit.
  - simpl; intro q; simpl.
    apply equal_refl.
Qed.
Next Obligation.
  intros m mon X Y Z W f; simpl; intro p; simpl.
  apply equal_sym.
  eapply equal_trans; [ apply bind_unit |]; simpl.
  simpl.
  apply equal_sym.
  eapply equal_trans; [ apply bind_assoc |]; simpl.
  eapply equal_trans; [ apply bind_assoc |]; simpl.
  apply eqMapEq, eqSubst.
  apply equal_sym.
  eapply equal_trans; [| apply compA ].
  apply equal_sym.
  eapply equal_trans; [ apply compMapEq |].
  - apply equal_refl.
  - simpl; intro q; simpl; apply bind_unit.
  - simpl; intro q; simpl.
    eapply equal_trans; [ apply bind_unit |].
    simpl.
    apply equal_refl.
Qed.
Next Obligation.
  intros m mon X Y Z k; simpl; intro p; simpl.
  eapply equal_trans; [ apply bind_unit |].
  simpl.    
  apply equal_refl.
Qed.
Next Obligation.
  intros m mon X Y Z W f g; simpl; intro p; simpl.
  eapply equal_trans; [ apply bind_assoc |]; simpl.
  apply equal_sym.
  eapply equal_trans; [ apply bind_assoc |]; simpl.
  apply equal_sym.
  apply eqMapEq, eqSubst; simpl; intro y; simpl.
  apply equal_sym.
  eapply equal_trans; [ apply bind_unit |].
  simpl.
  apply equal_refl.
Qed.
(* Program Instance kleisliArrow `(mon: Monad m): Arrow (Kleisli m) := *)
(*   { pure X Y := makeMap (fun f => unit \o f); *)
(*     conn X Y Z := makeMap (fun f => makeMap (fun g => bind g \o f)); *)
(*     frst X Y Z := makeMap (fun f: X --> m Y => makeMap (fun (p: prodSetoid X Z) => (f (fst p) >>= makeMap (fun q => unit ((q, snd p): prodSetoid Y Z))))) *)
(*   }. *)


Program Definition cokleisliArrow `(com: Comonad w): Arrow (CoKleisli w).
Abort.