This file implements bimonads. A Bimonad is a functor equipped with a monad/comonad structure and a distributive law over itself, which ensures the composition (W ∘ W) is also a monad and a comonad. Additionally there are some coherence laws, below, which may be read into several equivalent ways:
• The cojoin is a monad homomorphism (W → W ∘ W) and the counit a homomorphism (W → 1)
• The join is a comonad homomorphism (W ∘ W → W) and the unit a homomorphism (1 → W)
• W is a monoid in the category of comonoids (in the category of endofunctors)
• W is a comonoid in the category of monoids (in the category of endofunctors)

From Tealeaves.Classes Require Export

Import Product.Notations.
Import Functor.Notations.
#[local] Open Scope tealeaves_scope.

Context
(W : Type -> Type)
`{Fmap W}
`{Return W}
`{Join W}
`{Cojoin W}
`{Extract W}
`{BeckDistributiveLaw W W}.

`(fmap W (extract W) bdist W W = extract W (A := W A));
`(extract W bdist W W = fmap W (extract W (A := A)));
`(extract W join W = extract W extract W (A := W A));
`(extract W ret W = @id A);
`(cojoin W ret W = ret W ret W (A := A));
`(cojoin W join W (A := A) =
fmap W (join W) join W fmap W (bdist W W) cojoin W fmap W (cojoin W));
}.

# Kleisli lifting operation `bibind`

Context
(W : Type -> Type)
`{Fmap W} `{Join W} `{Cojoin W}
`{BeckDistribution W W}.

Definition bibind {A B} : (W A -> W B) -> (W A -> W B) :=
fun f => join W fmap W f cojoin W.

Definition twist : W W W W
:= fun A => fmap W (join W) bdist W W fmap W (cojoin W).

Definition kcomposebi {A B C} :
(W B -> W C) ->
(W A -> W B) ->
(W A -> W C) :=
fun g f => join W fmap W g twist B fmap W f cojoin W.

Module Notations.
Notation "g ⋆bi f" := (kcomposebi _ g f) (at level 70) : tealeaves_scope.
End Notations.

Import Notations.

## Specifications for suboperations

Context

Lemma cobind_to_bibind : forall `(f : W A -> B),
cobind W f = bibind W (ret W f).
Proof.
intros. unfold bibind. unfold_ops @Cobind_Cojoin.
rewrite <- (fun_fmap_fmap W). reassociate <-.
now rewrite (mon_join_fmap_ret W).
Qed.

Lemma bind_to_bibind : forall `(f : A -> W B),
bind W f = bibind W (f extract W).
Proof.
intros. unfold bibind. unfold_ops @Bind_Join.
rewrite <- (fun_fmap_fmap W).
reassociate <-. reassociate ->.
now rewrite (com_fmap_extr_cojoin W).
Qed.

Lemma fmap_to_bibind : forall `(f : A -> B),
fmap W f = bibind W (ret W f extract W).
Proof.
intros. unfold bibind.
do 2 rewrite <- (fun_fmap_fmap W).
repeat reassociate <-.
rewrite (mon_join_fmap_ret W).
reassociate ->.
now rewrite (com_fmap_extr_cojoin W).
Qed.

Context

Composition where `g` is a co-Kleisli arrow reduces to usual Kleisli composition.
Lemma bi_kcompose1 {A B C} : forall (g : B -> W C) (f : W A -> W B),
(g extract W) bi f = g f.
Proof.
intros. unfold kcomposebi, kcompose.
rewrite <- (fun_fmap_fmap W).
reassociate <-. unfold twist.
repeat reassociate <-.
reassociate -> near (fmap W (join W)).
rewrite (fun_fmap_fmap W). rewrite (bimonad_cap W).
rewrite <- (fun_fmap_fmap W).
repeat reassociate <-.
reassociate -> near (bdist W W).
reassociate -> near (fmap W (cojoin W)).
rewrite <- (natural (ϕ := @extract W _)); unfold_ops @Fmap_I.
reassociate <-. reassociate -> near (cojoin W (A := B)).
rewrite (com_fmap_extr_cojoin W).
reassociate -> near (fmap W f).
rewrite <- (natural (ϕ := @extract W _)); unfold_ops @Fmap_I.
repeat reassociate ->. rewrite (com_extract_cojoin W).
reflexivity.
Qed.

Composition where `f` is a Kleisli arrow reduces to usual co-Kleisli composition.
Lemma bi_kcompose2 {A B C} : forall (g : W B -> W C) (f : W A -> B),
g bi (ret W f) = g co f.
Proof.
intros. unfold kcomposebi, kcompose.
rewrite <- (fun_fmap_fmap W).
reassociate <-. unfold twist.
repeat reassociate <-.
reassociate -> near (fmap W (ret W)).
rewrite (fun_fmap_fmap W). rewrite (bimonad_cup W).
rewrite <- (fun_fmap_fmap W).
repeat reassociate <-.
reassociate -> near (fmap W (ret W) (A := W B)).
rewrite (dist_unit_r W W).
reassociate -> near (ret W).
rewrite (natural (ϕ := @ret W _)); unfold_ops @Fmap_I.
reassociate <-. reassociate -> near (fmap W (ret W)).
rewrite (mon_join_fmap_ret W).
reassociate -> near (ret W).
rewrite (natural (ϕ := @ret W _)); unfold_ops @Fmap_I.
reassociate <-. rewrite (mon_join_ret W).
reflexivity.
Qed.

## Functoriality of bibind

Context

Definition bind_id {A} :
bibind W (ret W extract W) = @id (W A).
Proof.
intros. unfold bibind.
rewrite <- (fun_fmap_fmap W).
do 2 reassociate -> on left.
rewrite (com_fmap_extr_cojoin W).
reassociate <- on left.
now rewrite (mon_join_fmap_ret W).
Qed.

Definition bind_functorial {A B C} : forall (g : W B -> W C) (f : W A -> W B),
bibind W (g bi f) = bibind W g bibind W f.
Proof.
intros. unfold bibind. unfold kcomposebi.
repeat reassociate -> on left.
rewrite <- (fun_fmap_fmap W).
repeat reassociate <- on left.
rewrite <- (mon_join_join W).
rewrite <- (fun_fmap_fmap W).
repeat reassociate -> on left.
rewrite <- (com_cojoin_cojoin W).
repeat reassociate <- on left.
repeat reassociate -> on left.
rewrite <- (fun_fmap_fmap W).
repeat reassociate <- on left.
reassociate -> near (fmap W (fmap W g)).
change (fmap W (fmap W g)) with (fmap (W W) g).
Set Keyed Unification.
rewrite <- (natural (ϕ := @join W _)).
Unset Keyed Unification.
rewrite <- (fun_fmap_fmap W).
repeat reassociate -> on left.
reassociate <- near (fmap W (fmap W f)).
change (fmap W (fmap W f)) with (fmap (W W) f).
Set Keyed Unification.
rewrite (natural (ϕ := @cojoin W _)).
Unset Keyed Unification.
repeat reassociate <- on left.
repeat reassociate <- on right.
reassociate -> near (join W).
repeat reassociate <-.
reassociate -> near (join W).
Set Keyed Unification.
rewrite (natural (ϕ := @join W _) (μ W)).
Unset Keyed Unification.
reassociate -> near (fmap W (cojoin W)).
Set Keyed Unification.
rewrite <- (natural (ϕ := @cojoin W _) (cojoin W)).
Unset Keyed Unification.
unfold twist.
repeat reassociate <-.
now repeat rewrite <- (fun_fmap_fmap W).
Qed.