Category-indexed monads, now with more type inference!

CatMonad.hs

-- https://stackoverflow.com/a/57046042/477476
-- https://gist.github.com/Lysxia/7331df3abee0aceacd3d9a74b567f54c

{-# LANGUAGE RankNTypes, TypeFamilies, PolyKinds, DataKinds #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs, TypeOperators #-}
{-# LANGUAGE TypeApplications, RebindableSyntax #-}

import Prelude hiding (id, (.))
import Control.Category
import Data.Kind
import GHC.TypeLits

class CatMonad (c :: k -> k -> Type) (m :: forall (x :: k) (y :: k). c x y -> Type -> Type) | c -> m where
  type Id c :: c x x
  type Cat c (f :: c x y) (g :: c y z) :: c x z

  xpure :: a -> m (Id c) a
  xbind :: m f a -> (a -> m g b) -> m (Cat c f g) b

-- Example #1: Indexed writer for an arbitrary Category.
-- There can be edges between any two objects.
data CatEdge (cat :: k -> k -> Type) (x :: k) (y :: k) = CE

newtype IWriter (cat :: k -> k -> Type) :: forall (i :: k) (j :: k). CatEdge cat i j -> Type -> Type where
  IWriter :: { runIWriter :: (a, cat i j) } -> IWriter cat (e :: CatEdge cat i j) a

instance Category cat => CatMonad (CatEdge cat) (IWriter cat) where
  type Id (CatEdge cat) = CE
  type Cat (CatEdge cat) _ _ = CE

  xpure a = IWriter (a, id)
  xbind (IWriter (a, f)) k =
    let IWriter (b, g) = k a in
    IWriter (b, f >>> g)

itell :: (Category cat) => cat i j -> IWriter cat (CE :: CatEdge cat i j) ()
itell f = IWriter ((), f)

ilisten :: (Category cat) => IWriter cat (CE :: CatEdge cat i j) a -> IWriter cat (CE :: CatEdge cat i j) (a, cat i j)
ilisten w = IWriter $
    let (x, f) = runIWriter w
    in ((x, f), f)

ipass :: (Category cat) => IWriter cat (CE :: CatEdge cat i j) (a, cat i j -> cat i j) -> IWriter cat (CE :: CatEdge cat i j) a
ipass w = IWriter $
    let ((x, censor), f) = runIWriter w
    in (x, censor f)

helloWriter :: IWriter (->) (CE :: CatEdge (->) Double Bool) String
helloWriter = do
    itell round
    itell even
    return True
    itell not
    return "foo"
  where
    -- Note: these MUST be defined non-pointfree, otherwise their type
    -- is not generalized enough.
    pure x = xpure x
    m >>= n = xbind m n

    return x = pure x
    m >> n = m >>= const n

-- Example #2: Counter. Index category is the (Nat, 0, +) monoid: a
-- single object, and one morphism per natural number.
data NatEdge (x :: ()) (y :: ()) where
    NE :: Nat -> NatEdge '() '()

type family NatPlus (n :: NatEdge x y) (m :: NatEdge y z) :: NatEdge x z where
    NatPlus (NE n) (NE m) = NE (n + m)

newtype Counter :: forall (i :: ()) (j :: ()). NatEdge i j -> Type -> Type where
    Counter :: { runCounter :: a } -> Counter (e :: NatEdge i j) a

instance CatMonad NatEdge Counter where
    type Id NatEdge = NE 0
    type Cat NatEdge f g = NatPlus f g

    xpure = Counter
    xbind act k = Counter $ runCounter $ k $ runCounter act

tick :: Counter (NE 1) ()
tick = Counter ()

-- > :t hello
-- hello :: Counter ('NE 3) Integer
helloCounter = do
    tick
    tick
    x <- pure 4
    tick
    y <- pure 5
    pure $ x + y
  where
    pure x = xpure x
    m >>= n = xbind m n

    return x = pure x
    m >> n = m >>= const n