Yoneda.hs

{-# OPTIONS_GHC -Wno-orphans #-}

-- | The Yoneda construction: @'Yoneda' p@ is the cofree profunctor on an arbitrary type of kind
-- @j +-> k@ (the 'HasCofree' instance for 'Profunctor'), and 'Yo' is the Yoneda embedding. By the
-- Yoneda lemma @Yoneda p@ is equivalent to @p@ when @p@ is already a profunctor
-- ('yoneda'\/'mkYoneda').
module Proarrow.Profunctor.Instance.Yoneda where

import Data.Function (($))
import Prelude ((*))

import Proarrow.Category.Enriched.Finitary (Finitary (..), FiniteCat, pairIndex, unpairIndex)
import Proarrow.Category.Instance.Nat (Nat (..))
import Proarrow.Category.Instance.Opposite (OPPOSITE (..), Op (..))
import Proarrow.Category.Instance.Prof (Prof (Prof))
import Proarrow.Core (CategoryOf (..), Hom, Profunctor (..), Promonad (..), (//), (:~>), type (+->))
import Proarrow.Functor (Functor (..))
import Proarrow.Profunctor.Cofree (HasCofree (..))
import Proarrow.Profunctor.Instance.Costar (Costar, pattern Costar)
import Proarrow.Profunctor.Instance.Star (Star, pattern Star)

-- | The cofree profunctor on @p@ (the 'HasCofree' instance for 'Profunctor'): natural
-- transformations out of the Yoneda embedding 'Yo'. Equivalent to @p@ when @p@ is already a
-- profunctor ('yoneda'\/'mkYoneda').
type Yoneda :: (j +-> k) -> j +-> k
data Yoneda p a b where
  Yoneda :: (Ob a, Ob b) => {unYoneda :: Yo a (OP b) :~> p} -> Yoneda p a b

instance (CategoryOf j, CategoryOf k) => Profunctor (Yoneda (p :: j +-> k)) where
  dimap l r (Yoneda k) = l // r // Yoneda \(Yo ca bd) -> k $ Yo (l . ca) (bd . r)
  r \\ Yoneda{} = r

instance Functor Yoneda where
  map (Prof n) = Prof \(Yoneda k) -> Yoneda (n . k)

instance Promonad (Star Yoneda) where
  id = Star (Prof mkYoneda)
  Star (Prof l) . Star (Prof r) = Star (Prof (l . yoneda . r))

instance Promonad (Costar Yoneda) where
  id = Costar (Prof yoneda)
  Costar (Prof l) . Costar (Prof r) = Costar (Prof (l . mkYoneda . r))

instance HasCofree Profunctor where
  type Cofree Profunctor p = Yoneda p
  lower = Prof yoneda
  unfoldMap (Prof n) = Prof (mkYoneda . n)

yoneda :: (CategoryOf j, CategoryOf k) => Yoneda (p :: j +-> k) :~> p
yoneda (Yoneda k) = k $ Yo id id

mkYoneda :: (Profunctor p) => p :~> Yoneda p
mkYoneda p = p // Yoneda \(Yo ca bd) -> dimap ca bd p

-- | Yoneda embedding
type Yo :: k -> OPPOSITE j -> j +-> k
data Yo a b c d where
  Yo :: c ~> a -> b ~> d -> Yo a (OP b) c d

instance (CategoryOf j, CategoryOf k) => Profunctor (Yo (a :: k) (OP b :: OPPOSITE j) :: j +-> k) where
  dimap l r (Yo f g) = Yo (f . l) (r . g)
  r \\ Yo f g = r \\ f \\ g

-- | The embedding is finitary when the arrows are: its elements over @c@\/@d@ are an arrow @c ~> a@
-- paired with an arrow @b ~> d@, numbered with the first varying slowest. This is the weight of the
-- ends in "Proarrow.Category.Enriched.Finitary.Topos", so it shares that module\'s
-- 'pairIndex' instead of spelling the radix out again.
instance (FiniteCat j, FiniteCat k, Ob a, Ob b) => Finitary (Yo (a :: k) (OP (b :: j)) :: j +-> k) where
  size @c @d = size @(Hom k) @c @a * size @(Hom j) @b @d
  toIndex @c @d (Yo ca bd) = pairIndex (size @(Hom j) @b @d) (toIndex @(Hom k) @c @a ca) (toIndex @(Hom j) @b @d bd)
  fromIndex @c @d i =
    let (l, r) = unpairIndex (size @(Hom j) @b @d) i
    in Yo (fromIndex @(Hom k) @c @a l) (fromIndex @(Hom j) @b @d r)

  -- spelled out for the same reason as the product's: the default would ask for a size per element
  elements @c @d = [Yo ca bd | ca <- elements @(Hom k) @c @a, bd <- elements @(Hom j) @b @d]

instance (CategoryOf j, CategoryOf k) => Functor (Yo (a :: k) :: OPPOSITE j -> j +-> k) where
  map (Op f) = Prof \(Yo ca bd) -> Yo ca (bd . f)

instance (CategoryOf j, CategoryOf k) => Functor (Yo :: k -> OPPOSITE j -> j +-> k) where
  map f = Nat (Prof \(Yo ca bd) -> Yo (f . ca) bd)