Free.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# OPTIONS_GHC -Wno-orphans #-}

-- | Free constructions: 'HasFree' captures the object constraints @ob@ whose forgetful functor has a left
-- adjoint, with @Free ob@ the free object, 'lift' the unit and 'foldMap' the universal property of the
-- adjunction (packaged as a 'Corepresentable' heteromorphism profunctor).
module Proarrow.Profunctor.Free where

import Data.Foldable1 (Foldable1 (foldMap1))
import Data.Kind (Constraint, Type)
import Data.List.NonEmpty (NonEmpty (..))
import Data.Maybe (Maybe (..))
import Prelude (($))
import Prelude qualified as P

import Proarrow.Category.Instance.Free (FREE (..), IsFreeOb (..), liftFree, retractFree)
import Proarrow.Category.Instance.IntConstruction (INT (..), IntConstruction (..), toInt)
import Proarrow.Category.Instance.Nat (Nat (..), first)
import Proarrow.Category.Instance.Prof (Prof (..))
import Proarrow.Category.Instance.Sub (Forget, On, SUBCAT (..), Sub (..))
import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), swap)
import Proarrow.Category.Monoidal.Applicative (Applicative (..))
import Proarrow.Category.Monoidal.CompactClosed (CompactClosed (..))
import Proarrow.Category.Monoidal.StarAutonomous (Dual, dualObj)
import Proarrow.Category.Monoidal.Strength (TracedMonoidal)
import Proarrow.Category.Monoidal.Strictified (Fold, Strictified (..), (==))
import Proarrow.Core
  ( CAT
  , CategoryOf (..)
  , Hom
  , Kind
  , OB
  , Profunctor (..)
  , Promonad (..)
  , UN
  , arr
  , lmap
  , obj
  , rmap
  , tgt
  , (//)
  , (:~>)
  )
import Proarrow.Functor (Functor (..))
import Proarrow.Limit.BinaryProduct (HasBinaryProducts)
import Proarrow.Limit.Terminal (HasTerminalObject)
import Proarrow.Monoid (Monoid)
import Proarrow.Monoid qualified as M
import Proarrow.Profunctor.Corepresentable (Corepresentable (..))
import Proarrow.Profunctor.Instance.Composition ((:.:) (..))
import Proarrow.Profunctor.Instance.Identity (Id)
import Proarrow.Profunctor.Instance.List (LIST (..), List (..))
import Proarrow.Profunctor.Instance.Star (Star, pattern Star)
import Proarrow.Profunctor.Representable (Rep (..))

type HasFree :: forall {k}. OB k -> Constraint
class (CategoryOf k, forall a. (Ob a) => ob (Free ob a)) => HasFree (ob :: OB k) where
  type Free ob (a :: k) :: k
  lift :: (Ob a) => a ~> Free ob a
  foldMap :: (ob b) => (a ~> b) -> Free ob a ~> b

retract :: forall ob a. (HasFree ob, ob a, Ob a) => Free ob a ~> a
retract = foldMap @ob id

freeMap :: (HasFree ob) => (a ~> b) -> Free ob a ~> Free ob b
freeMap @ob f = foldMap @ob (lift @ob . f) \\ f

freeComp :: (HasFree ob, Ob c) => b ~> Free ob c -> a ~> Free ob b -> a ~> Free ob c
freeComp @ob l r = foldMap @ob l . r

-- | By creating the left adjoint to the forgetful functor,
-- we obtain the free-forgetful adjunction.
instance (HasFree ob) => Corepresentable (Rep (Forget (ob :: OB k))) where
  type Rep (Forget ob) %% a = SUB (Free ob a)
  coindex (Rep f) = Sub (foldMap @ob f) \\ f
  corepUniv @a = let f = lift @ob @a in Rep f \\ f

instance HasFree P.Monoid where
  type Free P.Monoid a = [a]
  lift = P.pure
  foldMap = P.foldMap

instance HasFree P.Semigroup where
  type Free P.Semigroup a = NonEmpty a
  lift = P.pure
  foldMap = foldMap1

instance HasFree (P.Monoid `On` P.Semigroup) where
  type Free (P.Monoid `On` P.Semigroup) (SUB a) = SUB (Maybe a)
  lift = Sub Just
  foldMap (Sub f) = Sub (P.foldMap f)

-- | The free 'Applicative' on a functor @f@ (the 'HasFree' instance for 'Applicative'): formal 'pure',
-- effect and 'liftA2' nodes, retracted into any applicative by 'retractAp'.
type Ap :: (k -> Type) -> k -> Type
data Ap f a where
  Pure :: Unit ~> a -> Ap f a
  Eff :: f a -> Ap f a
  LiftA2 :: (Ob a, Ob b) => (a ** b ~> c) -> Ap f a -> Ap f b -> Ap f c

instance (CategoryOf k, Functor f) => Functor (Ap (f :: k -> Type)) where
  map f (Pure a) = Pure (f . a)
  map f (Eff x) = Eff (map f x)
  map f (LiftA2 k x y) = LiftA2 (f . k) x y

instance Functor Ap where
  map (Nat n) = Nat $ \case
    Pure a -> Pure a
    Eff fa -> Eff (n fa)
    LiftA2 k x y -> LiftA2 k (first (Nat n) x) (first (Nat n) y)

instance (Monoidal k) => Promonad (Star Ap :: CAT (k -> Type)) where
  id = Star (lift @Applicative)
  Star l . Star r = Star (freeComp @Applicative l r)

instance (Monoidal k, Functor f) => Applicative (Ap (f :: k -> Type)) where
  pure a () = Pure a
  liftA2 f (fa, fb) = LiftA2 f fa fb

-- | Given as 'P.Semigroup'\/'P.Monoid' rather than as 'Monoid' directly: @'Ap' f m@ is of kind
-- @Type@, where 'Monoid' already comes from the blanket @'P.Monoid' m => 'Monoid' (m :: Type)@
-- instance, so defining it here too would make every use overlap and solve to neither.
instance (Monoidal k, Monoid m) => P.Semigroup (Ap (f :: k -> Type) m) where
  l <> r = LiftA2 M.mappend l r

instance (Monoidal k, Monoid m) => P.Monoid (Ap (f :: k -> Type) m) where
  mempty = Pure M.mempty

retractAp :: (Applicative f) => Ap f a -> f a
retractAp (Pure a) = pure a ()
retractAp (Eff fa) = fa
retractAp (LiftA2 k x y) = liftA2 k (retractAp x, retractAp y)

instance (Monoidal k) => HasFree (Applicative :: OB (k -> Type)) where
  type Free Applicative f = Ap f
  lift = Nat Eff
  foldMap f@Nat{} = Nat retractAp . map f

-- | The free 'Promonad' on a profunctor @p@: a chain of @p@s ending in a hom arrow, folded into any
-- promonad by 'foldFreePromonad'.
data FreePromonad p a b where
  Unit :: (a ~> b) -> FreePromonad p a b
  Comp :: p a b -> FreePromonad p b c -> FreePromonad p a c

freePromonadAlg :: p :.: FreePromonad p :~> FreePromonad p
freePromonadAlg (p :.: pp) = Comp p pp

foldFreePromonad :: (Promonad q) => p :~> q -> FreePromonad p :~> q
foldFreePromonad _ (Unit f) = arr f
foldFreePromonad n (Comp p pp) = foldFreePromonad n pp . n p

instance (Profunctor p) => Profunctor (FreePromonad p) where
  dimap l r (Unit f) = Unit (r . f . l)
  dimap l r (Comp p q) = Comp (lmap l p) (rmap r q)
  r \\ (Unit f) = r \\ f
  r \\ Comp p q = r \\ p \\ q

instance (Profunctor p) => Promonad (FreePromonad p) where
  id = Unit id
  p . Unit f = lmap f p
  p . Comp r q = Comp r (p . q)
instance Functor FreePromonad where
  map = freeMap @Promonad
instance Promonad (Star FreePromonad) where
  id = Star (lift @Promonad)
  Star l . Star r = Star (freeComp @Promonad l r)
instance HasFree Promonad where
  type Free Promonad p = FreePromonad p
  lift = Prof \p -> p `Comp` Unit (tgt p)
  foldMap (Prof n) = Prof (foldFreePromonad n)

-- | The free @c@-structured kind over a @b@-structured kind @k@. A standalone family (rather than
-- an associated type of 'HasFreeK') because the kinds of 'Lift' and 'Retract' mention it.
type family FreeK (b :: Kind -> Constraint) (c :: Kind -> Constraint) (k :: Kind) :: Kind

-- | 'Proarrow.Object.Ob''-style helper: the quantified superclass of 'HasFreeK' needs to state
-- @c ('FreeK' b c k)@, and a type family application cannot head a quantified constraint directly.
class (c (FreeK b c k)) => FreeK' (b :: Kind -> Constraint) (c :: Kind -> Constraint) (k :: Kind)

instance (c (FreeK b c k)) => FreeK' b c k

-- | The embedding of an object of @k@ into the free @c@-structured kind.
type family Lift (b :: Kind -> Constraint) (c :: Kind -> Constraint) (a :: k) :: FreeK b c k

-- | Interpret an object of the free @c@-structured kind back into @k@.
type family Retract (b :: Kind -> Constraint) (c :: Kind -> Constraint) (k :: Kind) (a :: FreeK b c k) :: k

-- | @'FreeK' b c@ builds the free @c@-structured kind over any @b@-structured kind: 'liftK'
-- embeds the arrows of @k@, and when @k@ itself is already @c@-structured 'retractK' interprets
-- back into @k@.
class
  (forall k. (b k) => FreeK' b c k) =>
  HasFreeK (b :: Kind -> Constraint) (c :: Kind -> Constraint)
  where
  liftK :: (b k) => (x :: k) ~> y -> Lift b c x ~> Lift b c y
  retractK
    :: forall k (x :: FreeK b c k) (y :: FreeK b c k)
     . (c k)
    => x ~> y -> Retract b c k x ~> Retract b c k y

type instance FreeK CategoryOf Monoidal k = LIST k

type instance Lift CategoryOf Monoidal a = L '[a]
type instance Retract CategoryOf Monoidal k (a :: LIST k) = Fold (UN L a)

instance HasFreeK CategoryOf Monoidal where
  liftK f = Cons f Nil
  retractK Nil = one
  retractK (Cons f Nil) = f
  retractK (Cons f fs@Cons{}) = f ** retractK @CategoryOf @Monoidal fs

-- | The free category with a terminal object over @k@, built with
-- "Proarrow.Category.Instance.Free".
type instance FreeK CategoryOf HasTerminalObject k = FREE '[HasTerminalObject] (Hom k)

type instance Lift CategoryOf HasTerminalObject (a :: k) = EMB a
type instance Retract CategoryOf HasTerminalObject k (a :: FREE '[HasTerminalObject] (Hom k)) = Lower (Id :: CAT k) a

instance HasFreeK CategoryOf HasTerminalObject where
  liftK = liftFree
  retractK = retractFree @'[HasTerminalObject]

-- | The free category with binary products over @k@, built with
-- "Proarrow.Category.Instance.Free".
type instance FreeK CategoryOf HasBinaryProducts k = FREE '[HasBinaryProducts] (Hom k)

type instance Lift CategoryOf HasBinaryProducts (a :: k) = EMB a
type instance Retract CategoryOf HasBinaryProducts k (a :: FREE '[HasBinaryProducts] (Hom k)) = Lower (Id :: CAT k) a

instance HasFreeK CategoryOf HasBinaryProducts where
  liftK = liftFree
  retractK = retractFree @'[HasBinaryProducts]

type instance FreeK TracedMonoidal CompactClosed k = INT k

type instance Lift TracedMonoidal CompactClosed (a :: k) = I a Unit
type instance Retract TracedMonoidal CompactClosed k (I a b :: INT k) = a ** Dual b

instance HasFreeK TracedMonoidal CompactClosed where
  liftK = toInt
  retractK (Int @ap @am @bp @bm f) =
    dualObj @am //
      dualObj @bm //
        unStr $
          Str @[ap, Dual am] @[Dual am, ap] (swap @_ @ap @(Dual am)) ** Str @'[] @[bm, Dual bm] (dualityUnit @_ @bm)
            == obj @'[Dual am] ** Str @[ap, bm] @[am, bp] f ** obj @'[Dual bm]
            == Str @[Dual am, am] @'[] (dualityCounit @_ @am) ** obj @'[bp] ** obj @'[Dual bm]