Category-indexed monads

CatMonad.hs

-- https://stackoverflow.com/a/57046042/477476
-- https://gist.github.com/Lysxia/04039e4ca6f7a3740281e4e3583ae971

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

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

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

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

data CatMonadKit (m :: forall (x :: k) (y :: k). c x y -> Type -> Type) = CatMonadKit
    { pure, return :: forall a. a -> m (Id m) a
    , (>>=) :: forall f g a b. m f a -> (a -> m g b) -> m (Cat m f g) b
    , (>>) :: forall f g a b. m f a -> m g b -> m (Cat m f g) b
    }

catMonad :: forall (m :: forall (x :: k) (y :: k). c x y -> Type -> Type). (CatMonad m) => CatMonadKit m
catMonad = CatMonadKit
    { pure = xpure
    , return = xpure
    , (>>=) = xbind
    , (>>) = \m n -> xbind m (const n)
    }

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

newtype IWriter cat i j (q :: CatEdge i j) a = IWriter { runIWriter :: (a, cat i j) }

instance Category cat => CatMonad (IWriter cat) where
  type Id (IWriter cat) = CE
  type Cat (IWriter 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 i j CE ()
itell f = IWriter ((), f)

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

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

helloWriter :: IWriter (->) Double Bool CE String
helloWriter = do
    itell round
    itell even
    return True
    itell not
    return "foo"
  where
    -- doesn't quite work yet :/
    -- CatMonadKit{..} = catMonad @_ @_ @(IWriter (->))

    -- Note: these MUST be defined non-pointfree, otherwise their type
    -- is not generalized enough.
    pure x = xpure @_ @_ @(IWriter (->)) x
    m >>= n = xbind @_ @_ @(IWriter (->)) 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 (i :: ()) (j :: ()) (q :: NatEdge i j) a = Counter { runCounter :: a }

instance CatMonad Counter where
    type Id Counter = NE 0
    type Cat Counter 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 '() '() ('E 3) Integer
helloCounter = do
    tick
    tick
    x <- pure 4
    tick
    y <- pure 5
    pure $ x + y
  where
    -- CatMonadKit{..} = catMonad @_ @_ @Counter

    pure x = xpure @_ @_ @Counter x
    m >>= n = xbind @_ @_ @Counter m n

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