/Category.lean at fca8be10990480f321cb38647a473578d3a53da3

import Mathlib

/-
Lean, the category!
Objects: Types
Morphisms: (Total) Functions

Resources:
https://raw.githubusercontent.com/BartoszMilewski/DaoFP/refs/heads/master/DaoFP.pdf
https://math.andrej.com/2016/08/06/hask-is-not-a-category/
https://www.mit.edu/~xy/lean/
https://tannerduve.github.io/files/monads.pdf

TODO: Diagrams?
-/

-- Initial object
#check Empty

-- Morphism from initial object
#check Empty.elim

-- Terminal object
#check PUnit

-- Identity morphism
#check id

-- Composition of morphism is associative
#check Function.comp_assoc

-- Equality of morphisms
#check funext

-- Sums (coproducts)
#check Sum

-- Products
#check Prod

-- Exponentials
#check (· → ·)

-- Lean is a bicartesian closed category, which means it has an initial object, terminal object, sums, products, exponentials, and sums distribute over products.


/-
Endofunctors in Lean!

f maps objects
f.map or `<$>` maps morphisms
`(α → β) → f α → f β` is the same thing as `(α → β) → (f α → f β)`
-/
#check Functor
#check LawfulFunctor

#synth Functor List

#synth Functor Option

#synth Functor Tree

#synth Functor (Except String)

instance : LawfulFunctor List where
  map_const := by solve_by_elim
  id_map xs := by simp
  comp_map := by simp


/--
Contravariant functors (normal functors are "covariant" functors) are functors from Colean to Lean

`map` turns a "producer of α" into a "producer of β"
`comap` turns a "consumer of α" into a "consumer of β"
-/
class Cofunctor (f : Type u → Type v) where
  comap : (β → α) → f α → f β
  id_comap (x : f α) : comap id x = x
  comp_comap (g : β → α) (h : γ → β) (x : f α) : comap (g ∘ h) x = comap h (comap g x)

/-- This is not standard notation but I just made something up -/
infixr:100 " <¥> " => Cofunctor.comap

theorem Cofunctor.id_comap' [Cofunctor f] : Cofunctor.comap (f := f) (@id α) = id := by
  ext x
  exact Cofunctor.id_comap x

theorem Cofunctor.comap_comp_comap [Cofunctor f] (g : α → β) (h : β → γ) :
    ((g <¥> ·) ∘ (h <¥> ·) : f γ → f α) = Cofunctor.comap (h ∘ g) :=
  funext fun _ ↦ (comp_comap _ _ _).symm


/-- Hom-functor in enriched category -/
@[simp]
instance (α : Type u) : Functor (α → ·) where
  map f g := f ∘ g

instance (α : Type u) : LawfulFunctor (α → ·) where
  map_const := by solve_by_elim
  id_map := by simp
  comp_map := by simp [Function.comp_assoc]

@[simp]
instance (α : Type u) : Cofunctor (· → α) where
  comap f g := g ∘ f
  id_comap := by simp
  comp_comap := by simp [Function.comp_assoc]

-- Most examples of cofunctors in Lean are these function object things

/-
A function type is covariant if the free param is in an even depth and contravariant otherwise.

α → · is co
· → α is contra
(· → α) → β is co
((· → α) → β) → γ is contra
and so on
-/


/-- Composition of two functors of same variance is a functor -/
@[simp]
instance [Functor f] [Functor g] : Functor (f ∘ g) where
  map h x := Functor.map (f := f) (h <$> ·) x

instance [Functor f] [LawfulFunctor f] [Functor g] [LawfulFunctor g] : LawfulFunctor (f ∘ g) where
  map_const := by solve_by_elim
  id_map := by simp
  comp_map h h' x := by simp; rfl

@[simp]
instance [Cofunctor f] [Cofunctor g] : Functor (f ∘ g) where
  map h x := Cofunctor.comap (f := f) (Cofunctor.comap h) x

instance [Cofunctor f] [Cofunctor g] : LawfulFunctor (f ∘ g) where
  map_const := by solve_by_elim
  id_map := by simp [Cofunctor.id_comap']
  comp_map h h' x := by simp [← Cofunctor.comap_comp_comap]

-- If functors are sort of like "containers" for data, then functor composition is "nesting" two containers
#synth LawfulFunctor (List ∘ Option)


/-- Composition of functors of opposite variance is a contravariant functor -/
@[simp]
instance [Functor f] [LawfulFunctor f] [Cofunctor g] : Cofunctor (f ∘ g) where
  comap h x := Functor.map (f := f) (h <¥> ·) x
  id_comap := by simp [Cofunctor.id_comap]
  comp_comap := by simp [Cofunctor.comp_comap]

@[simp]
instance [Cofunctor f] [Functor g] [LawfulFunctor g] : Cofunctor (f ∘ g) where
  comap h x := Cofunctor.comap (f := f) (h <$> ·) x
  id_comap := by
    simp only [Function.comp_apply, id_map]
    exact Cofunctor.id_comap
  comp_comap h h' x := by
    simp only [← Functor.map_comp_map h' h, Cofunctor.comp_comap]


-- Bifunctors map Lean × Lean to Lean
#check Bifunctor
#check LawfulBifunctor

-- `Sum` and `Prod` are bifunctors
#synth LawfulBifunctor Sum
#synth LawfulBifunctor Prod


-- Multivariate functors
#check MvFunctor
#check LawfulMvFunctor


/--
Profunctors are useful for lenses and optics

For more info about optics see https://marcosh.github.io/post/2025/10/07/the-mondrian-introduction-to-functional-optics.html
-/
class Profunctor (p : Type u → Type v → Type*) where
  dimap : (s → a) → (b → t) → (p a b → p s t)
  id_dimap (x : p α β) : dimap id id x = x
  dimap_dimap (f : α₁ → α₀) (f' : α₂ → α₁) (g : β₀ → β₁) (g' : β₁ → β₂) (x : p α₀ β₀) :
    dimap f' g' (dimap f g x) = dimap (f ∘ f') (g' ∘ g) x

/-- Exponentials are profunctors -/
instance : Profunctor (· → ·) where
  dimap f g h := g ∘ h ∘ f
  id_dimap := by simp
  dimap_dimap := by simp [Function.comp_assoc]

inductive Procompose p q [Profunctor p] [Profunctor q] a b
  | mk : q a x → p x b → Procompose p q a b

def mapOut [Profunctor p] [Profunctor q] (pc : Procompose p q a b) (f : {x : Type u} → q a x → p x b → c) :=
  match pc with
  | ⟨qax, pxb⟩ => f qax pxb

instance [Profunctor p] [Profunctor q] : Profunctor (Procompose p q) where
  dimap l r
    | ⟨qax, pxb⟩ => ⟨Profunctor.dimap l id qax, Profunctor.dimap id r pxb⟩
  id_dimap := by simp [Profunctor.id_dimap]
  dimap_dimap := by simp [Profunctor.dimap_dimap]

abbrev End p [Profunctor p] := ∀ x, p x x

abbrev Coend p [Profunctor p] := Σ x, p x x

abbrev ProPair q p [Profunctor p] [Profunctor q] a b x y :=
  q a y × p x b

instance [Profunctor p] [Profunctor q] : Profunctor (ProPair q p a b) where
  dimap l r
    | ⟨qax, pxb⟩ => ⟨Profunctor.dimap id r qax, Profunctor.dimap l id pxb⟩
  id_dimap := by simp [Profunctor.id_dimap]
  dimap_dimap := by simp [Profunctor.dimap_dimap]

abbrev CoendCompose p q [Profunctor p] [Profunctor q] a b :=
  Coend (ProPair q p a b)

instance [Profunctor p] [Profunctor q] : Profunctor (CoendCompose p q) where
  dimap l r
    | ⟨x, (qay, pxb)⟩ => ⟨x, (Profunctor.dimap l id qay, Profunctor.dimap id r pxb)⟩
  id_dimap := by simp [Profunctor.id_dimap]
  dimap_dimap := by simp [Profunctor.dimap_dimap]

-- TODO: Existential lenses?

-- TODO: Tambara modules and profunctor optics?


/--
Type of a natural transformation (without the naturality condition)

Intuitively, it represents moving data from one "container" to another
-/
abbrev NaturalType.{u} (f : Type u → Type v) (g : Type u → Type v) :=
  {α : Type u} → f α → g α

/--
Naturality is automatically guarenteed for parametrically polymorphic functions (where the implementation is the same for each type), AKA "theorems for free". This is not guarenteed in general though since we could have a function which inspects the input type and does something crazy.

TODO: Would it be more convenient to make this a subtype?
-/
class Natural f [Functor f] [LawfulFunctor f] g [Functor g] [LawfulFunctor g] (η : NaturalType f g) where
  naturality (h : α → β) (x : f α) : h <$> (η x) = η (h <$> x)

instance : Natural List Option List.head? :=
  ⟨by simp⟩

def OptionToList : Option α → List α
  | some a => [a]
  | none => []

instance : Natural Option List OptionToList :=
  ⟨by simp [OptionToList]; grind⟩

class Conatural f [Cofunctor f] g [Cofunctor g] (η : NaturalType f g) where
  naturality (h : β → α) (x : f α) : h <¥> (η x) = η (h <¥> x)

/--
Vertical composition of natural transformations

Intuitively this is like doing two data moves.
-/
instance [Functor f] [LawfulFunctor f] [Functor g] [LawfulFunctor g] [Functor h] [LawfulFunctor h] [M : Natural f g η] [N : Natural g h μ] :
    Natural f h (fun {α : Type u} ↦ @μ α ∘ @η α) :=
  ⟨by simp [N.naturality, M.naturality]⟩

/--
Horizontal composition of natural transformations

Intuitively this is like repackaging data in nested "containers"
-/
instance (η : NaturalType f f') (μ : NaturalType g g') [Functor f] [LawfulFunctor f] [Functor f'] [LawfulFunctor f'] [Functor g] [LawfulFunctor g] [Functor g'] [LawfulFunctor g'] [M : Natural f f' η] [N : Natural g g' μ] :
    Natural (g ∘ f) (g' ∘ f') (μ ∘ (η <$> ·)) :=
  ⟨by simp [N.naturality, M.naturality]⟩

/-- Alternatively we do `μ` first and then the map second -/
instance (η : NaturalType f f') (μ : NaturalType g g') [Functor f] [LawfulFunctor f] [Functor f'] [LawfulFunctor f'] [Functor g] [LawfulFunctor g] [Functor g'] [LawfulFunctor g'] [M : Natural f f' η] [N : Natural g g' μ] :
    Natural (g ∘ f) (g' ∘ f') ((Functor.map (f := g') η ·) ∘ μ) :=
  ⟨by simp [N.naturality, M.naturality]⟩

/-- The two orderings are equivalent, which only requires the outer transformation to be natural -/
lemma horizontal_comp_equiv (η : NaturalType f f') (μ : NaturalType g g') [Functor f] [LawfulFunctor f] [Functor f'] [LawfulFunctor f'] [Functor g] [LawfulFunctor g] [Functor g'] [LawfulFunctor g'] [N : Natural g g' μ] :
    (μ ∘ (η <$> ·)) x = ((Functor.map (f := g') η ·) ∘ μ) x := by
  simp [N.naturality]


/-- Yoneda forward map (g is not necessarily natural) -/
def yoneda (g : NaturalType (α → ·) f) [Functor f] [LawfulFunctor f] : f α :=
  g id

/-- Yoneda reverse map -/
def yoneda' [Functor f] [LawfulFunctor f] (y : f α) : NaturalType (α → ·) f :=
  (· <$> y)

/-- Reverse map always produces a natural transformation -/
instance [Functor f] [LawfulFunctor f] : Natural (α → ·) f (yoneda' y) :=
  ⟨fun h x ↦ by simp [yoneda']; rfl⟩

/--
Mapping and unmapping a natural transformation returns the itself

Note that this does not work for an arbitrary function between the hom-functor and `f` because we use the naturality condition.
-/
theorem yoneda_lemma (g : NaturalType (α → ·) f) [Functor f] [LawfulFunctor f] [N : Natural (α → ·) f g] : yoneda' (yoneda g) x = g x := by
  simp [yoneda, yoneda', N.naturality]

/-- Mapping and unmapping an element `f α` returns itself -/
theorem yoneda_lemma' (y : f α) [Functor f] [LawfulFunctor f] : yoneda (yoneda' y) = y := by
  simp [yoneda, yoneda']

/-- Coyoneda forward map -/
def coyoneda (g : NaturalType (· → α) f) [Cofunctor f] : f α :=
  g id

/-- Coyoneda reverse map -/
def coyoneda' [Cofunctor f] (y : f α) : NaturalType (· → α) f :=
  (· <¥> y)

/-- Reverse map always produces a natural transformation -/
instance [Cofunctor f] : Conatural (· → α) f (coyoneda' y) :=
  ⟨fun h x ↦ by simp [coyoneda', Cofunctor.comp_comap]⟩

/-- Same but for Coyoneda -/
theorem coyoneda_lemma (g : NaturalType (· → α) f) [Cofunctor f] [N : Conatural (· → α) f g] : coyoneda' (coyoneda g) x = g x := by
  simp [coyoneda, coyoneda', N.naturality]

/-- Same but for Coyoneda -/
theorem coyoneda_lemma' (y : f α) [Cofunctor f] : coyoneda (coyoneda' y) = y := by
  simp [coyoneda, coyoneda', Cofunctor.id_comap]


-- Applicative functors
#check Applicative
#check LawfulApplicative

-- Motivation: mapping multi-arg functions
#simp (some 3).map (· * ·)

#eval (· * ·) <$> (some 3) <*> (some 4)

/-- Composition of two applicatives is an applicative -/
@[simp]
instance [Applicative f] [Applicative g] : Applicative (f ∘ g) where
  pure x := pure (f := f) (pure x)
  seq h x := Seq.seq (f := f) ((· <*> ·) <$> h) x

instance [Applicative f] [LawfulApplicative f] [Applicative g] [LawfulApplicative g] : LawfulApplicative (f ∘ g) where
  seqLeft_eq := by simp
  seqRight_eq := by simp
  pure_seq := by simp [pure_seq]
  map_pure := by simp
  seq_pure := by simp
  seq_assoc x h h' := by
    simp [seq_assoc, seq_map_assoc, map_seq]
    congr
    ext
    simp [seq_assoc]

-- TODO: Lax monoidal functors


-- Monads, "warm fuzzy things"
#check Monad
#check LawfulMonad

/-
Functors let us apply `α → β` to `f α`
Applicatives let us apply `f (a → β)` to `f α`
But what about applying an "effectful function" `α → f β` to `f α`?
Another use case is to compose `α → f β` and `β → f γ`.

Kleisli category: Any monad `m` creates a category where the objects are still types but the morphisms are `α → β` for every `α → f β` in Lean. Then composition of effectful functions becomes composition of morphisms.

This construction also motivates the monad laws.

In fact, using `>>=` and `pure` we can implement `<$>` and `<*>` so every monad is also a functor and applicative.

Exercise: Find an example of a functor which is not applicative and an applicative which is not a monad.
-/

#synth Monad Option

#synth Monad IO

#synth Monad (StateM ℕ)

#synth Monad (Writer ℕ)

#synth Monad (ST ℕ)

#synth Monad (Except String)

#synth Monad (Sum ℕ)

instance : LawfulMonad Option :=
  LawfulMonad.mk' Option
    (id_map := by simp)
    (pure_bind := by simp [Option.bind])
    (bind_assoc := by simp; grind)
    (bind_pure_comp := by simp [Option.map]; grind)

#synth Monad List

#synth LawfulMonad List


/--
This function looks ugly, but we can simplify it with `do` notation, which is syntactic sugar that lets us unwrap monadic values and automatically inserts `>>=` when we use the unwrapped values

https://slightknack.dev/blog/do-notation/
-/
def option_div (x_wrapped : Option ℕ) (y_wrapped : Option ℕ) : Option ℚ :=
  y_wrapped >>= fun y ↦
    if y = 0 then
      none
    else
      x_wrapped >>= fun x ↦ some <| x / y

#eval option_div (some 3) (some 0)

def option_div' (x_wrapped : Option ℕ) (y_wrapped : Option ℕ) : Option ℚ := do
  let x ← x_wrapped
  let y ← y_wrapped
  if y = 0 then none else some <| x / y

/-- Even the identity monad is powerful! -/
def Array.insSort [LinearOrder α] (A : Array α) := Id.run do
  let N := A.size
  let mut A := A.toVector
  for hi : i in [:N] do
    for hj : j in [:i] do
      have := Membership.get_elem_helper hi rfl
      if A[i - j] < A[i - j - 1] then
        A := A.swap (i - j - 1) (i - j)
      else
        break
  return A.toArray

/-- List monad demo -/
def UpToN (xs : List ℕ) : List ℕ := do
  let x ← xs
  let y ← List.range x
  return y

#eval UpToN [1, 2, 3]

/-
Sadly, in general monads do not compose 😿
https://carlo-hamalainen.net/2014/01/02/applicatives-compose-monads-do-not/

However, in some cases we can use monad transformers to compose them.
-/


-- Equivalent definition using "fish"
#check Bind.kleisliRight
-- Equivalent definition using "join"
#check joinM
-- Exercise: Implement bind using fish

/-
"A monad is just a monoid in the category of endofunctors"

In fact, there is a bijection between the two!

https://old.reddit.com/r/math/comments/ap25mr/a_monad_is_a_monoid_in_the_category_of/

The category of Lean endofunctors
Objects: Endofunctors
Morphisms: Natural transformations (we showed earlier that vertical composition produces another natural transformation)
-/

/-- Every object has an identity morphism -/
instance [Functor f] [LawfulFunctor f] : Natural f f id :=
  ⟨by simp⟩

/-- Vertical composition is associative -/
lemma nat_trans_comp_assoc (η : NaturalType f g) (μ : NaturalType g h) (ν : NaturalType h i) [Functor f] [LawfulFunctor f] [Functor g] [LawfulFunctor g] [Functor h] [LawfulFunctor h] [Functor i] [LawfulFunctor i] :
    ((ν ∘ μ) ∘ η) x = (ν ∘ μ ∘ η) x := by
  simp only [Function.comp_assoc]

#check Monoid

/-
A monoidal category is a category C equipped with a tensor product ⨂ from C × C to C and an identity object I with certain properties.

For the category of Lean endofunctors, let ⨂ be functor composition and I be the identity functor `Id`.
-/

/-- ⨂ is obviously associative -/
lemma functor_comp_assoc [Functor f] [LawfulFunctor f] [Functor g] [LawfulFunctor g] [Functor h] [LawfulFunctor h] : (f ∘ g) ∘ h = f ∘ g ∘ h := by
  apply Function.comp_assoc

/-- `Id` is an identity for ⨂ -/
lemma functor_left_id [Functor f] [LawfulFunctor f] : id ∘ f = f := by
  simp

/-- `Id` is an identity for ⨂ -/
lemma functor_right_id [Functor f] [LawfulFunctor f] : f ∘ id = f := by
  simp

-- The coherence conditions (insert scary pentagon diagram here) are automatically satisfied because the associator and unitor natural isomorphisms are equalities.

/-
A monoidal object is an object M in (C, ⨂, I) with an arrow μ from M ⨂ M to M and η from I to M such that μ is associative and η is an identity with respect to μ.

A monoidal object in the category of Lean endofunctors is a functor with natural transformations `join` (corresponding to μ) and `pure` (η) with the following properties:
-/

class EndofunctorMonoid m extends Functor m, LawfulFunctor m where
  join : NaturalType (m ∘ m) m
  pure : NaturalType Id m
  join_pure : (join ∘ pure) x = x
  -- When using <$>, Lean synthesizes the wrong type class instance for some weird reason
  join_map_pure : (join ∘ (map pure ·)) x = x
  join_join : (join ∘ (map join ·)) x = (join ∘ join) x

@[simp]
def bindFromJoin [EndofunctorMonoid m] (join : NaturalType (m ∘ m) m) (x : m α) (f : α → m β) :=
  join (Functor.map (f := m) f x)

@[simp]
instance [EndofunctorMonoid m] : Monad m where
  pure := EndofunctorMonoid.pure
  bind := bindFromJoin EndofunctorMonoid.join

/-- A monoid in the category of endofunctors is a monad -/
instance [EndofunctorMonoid m] [J : Natural (m ∘ m) m EndofunctorMonoid.join] [P : Natural Id m EndofunctorMonoid.pure] : LawfulMonad m :=
  LawfulMonad.mk' m id_map
    (pure_bind := fun x f ↦ by
      simpa [P.naturality, Functor.map] using EndofunctorMonoid.join_pure)
    (bind_assoc := fun x f g ↦ by
      have := EndofunctorMonoid.join_join (x := (fun a ↦ Functor.map (f := m) g (f a)) <$> x)
      simp at this
      simp [J.naturality, ← this])
    (map_const := by simp [map_const])
    (bind_pure_comp := fun f x ↦ by
      have := EndofunctorMonoid.join_map_pure (x := f <$> x)
      simpa using this)

@[simp]
def joinFromBind [Monad m] (bind : {α β : Type u} → m α → (α → m β) → m β) (x : m (m α)) :=
  bind x id

/-- A monad is a monoid in the category of endofunctors -/
@[simp]
instance [Monad m] [LawfulMonad m] : EndofunctorMonoid m where
  pure := pure
  join := joinFromBind bind
  join_pure := by simp
  join_map_pure := by simp
  join_join := by simp

instance [Monad m] [LawfulMonad m] : Natural (m ∘ m) m EndofunctorMonoid.join :=
  ⟨by simp⟩

instance [Monad m] [LawfulMonad m] : Natural Id m EndofunctorMonoid.pure :=
  ⟨by simp [Functor.map]⟩

/-- `bindFromJoin` and `joinFromBind` form a bijection -/
theorem bind_join_equiv [Monad m] [LawfulMonad m] : (bindFromJoin (m := m) (joinFromBind bind)) x f = bind x f := by
  simp

theorem bind_join_equiv' [E : EndofunctorMonoid m] : joinFromBind (bindFromJoin E.join) x = E.join x := by
  simp


-- TODO: Monad transformers


-- TODO: Enrichment


-- Unlike Haskell, Lean is powerful enough that we can also use it for doing category theory in any category, not just the category Lean
#check CategoryTheory.Category

#check CategoryTheory.Functor

#check CategoryTheory.yoneda

#check CategoryTheory.Monad

#check CategoryTheory.Monad.monadMonEquiv