monad homomorphisms and stuff

MonadHom.hs

{-# language ConstraintKinds #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language LambdaCase #-}
{-# language MultiParamTypeClasses #-}
{-# language PolyKinds #-}
{-# language RankNTypes #-}
{-# language ScopedTypeVariables #-}
{-# language TupleSections #-}
{-# language TypeFamilies #-}
{-# language TypeOperators #-}
{-# language UndecidableInstances #-}
{-# language UndecidableSuperClasses #-}
{-# options_ghc -Wall #-}
{-# options_ghc -fno-warn-orphans #-}
{-# options_ghc -fno-warn-deprecations #-}
{-# options_ghc -fwarn-redundant-constraints #-}

module Control.Monad.Homomorphism where

import Control.Applicative
import Control.Category
import Control.Monad.Codensity
import Control.Monad.Cont
import Control.Monad.Error
import Control.Monad.Except
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.RWS
import Control.Monad.State
import Control.Monad.Trans.Identity
import Control.Monad.Trans.Maybe
import Control.Monad.Writer
import Control.Monad.Zip
import Data.Constraint
import Data.Constraint.Lifting
import Data.Functor.Coyoneda
import Data.Functor.Yoneda
import Data.Proxy
import Data.Void
import Prelude hiding (id,(.))

--------------------------------------------------------------------------------
-- * monad homomorphisms
--
-- Common existing monad homomorphisms
--
-- @
-- lift           :: (MonadTrans t, Monad m) => m ~> t m -- monic unless t is terminal
-- liftIO         :: MonadIO m => IO ~> m -- monic unless m is terminal
-- ioToST         :: IO ~> ST RealWorld -- monad iso w/ inverse stToIO
-- stToIO         :: ST RealWorld ~> IO -- monad iso w/ inverse ioToST
-- unsafeIOToST   :: IO ~> GHC.ST.ST s -- unsafe monad iso w/ inverse unsafeSTToIO
-- unsafeSTToIO   :: GHC.ST.ST s ~> IO -- unsafe monad iso w/ inverse unsafeIOToST
-- lowerCodensity :: Monad m => Codensity m ~> m -- split epi w/ lift as right inverse
-- lowerYoneda    :: Yoneda m ~> m -- isomorphism w/ inverse lift
-- lowerCoyoneda  :: Functor m => Coyoneda m ~> m -- isomorphism w/ inverse lift
-- zoom l         -- monomorphism if l is a lens
-- magnify g      -- monomorphism if g is a getter
-- flip runReaderT :: e -> ReaderT e m ~> m -- epic
-- flip evalStateT :: Functor m => s -> StateT s m ~> m -- epic
-- @
--------------------------------------------------------------------------------

-- | monad homomorphism
--
-- @
-- f return = return
-- f (m >>= n) = f m >>= f . n
-- @
type m ~> n = forall a. m a -> n a

(<&>) :: Functor f => f a -> (a -> b) -> f b
(<&>) = flip fmap

transform :: Lifting Monad t => Monad m :- Monad (t m)
transform = lifting

-- | Note: ContT r is MPointed, but an MFunctor
class (Lifting Monad t, MonadTrans t) => MPointed t
instance (Lifting Monad t, MonadTrans t) => MPointed t

instance Lifting Monad Codensity where
  lifting = Sub Dict

instance Lifting Monad Yoneda where
  lifting = Sub Dict

instance Lifting Monad Coyoneda where
  lifting = Sub Dict

class Lifting Monad t => MFunctor t where
  -- @
  -- hoist id = id
  -- @
  --
  -- Given monad homomorphisms @f@ and @g@:
  --
  -- @
  -- hoist f . hoist g = hoist (f . g)
  -- @

  -- Given a monad homomorphism @f@, @hoist f@ is a monad homomorphism:
  --
  -- @
  -- hoist f (return a) = return a
  -- hoist f (m >>= n) = hoist f m >>= hoist f . n
  -- @
  --
  -- example:
  --
  -- @
  -- hoist init :: (MFunctor t, Monad m) => t Identity ~> t m
  -- @
  hoist :: (Monad m, Monad n) => (m ~> n) -> t m ~> t n

-- | @ContT r@ is a functor from the category of monad isomorphisms
-- over hask to the category of monads over hask
hoistContT :: (m ~> n) -> (n ~> m) -> ContT r m ~> ContT r n
hoistContT f g (ContT m) = ContT $ \k -> f (m (g . k))

-- | @Codensity@ is a functor from the category of monad isomorphisms (split monos?)
-- over hask to the category of monads over hask
hoistCodensity :: (m ~> n) -> (n ~> m) -> Codensity m ~> Codensity n
hoistCodensity f g (Codensity m) = Codensity $ \k -> f (m (g . k))
  
instance MFunctor IdentityT where
  hoist f (IdentityT m) = IdentityT (f m)

instance MFunctor (StateT s) where
  hoist f (StateT m) = StateT (f . m)

instance MFunctor (ReaderT e) where
  hoist f (ReaderT m) = ReaderT (f . m)

instance Monoid e => MFunctor (WriterT e) where
  hoist f (WriterT m) = WriterT (f m)

instance Monoid w => MFunctor (RWST r w s) where
  hoist f (RWST m) = RWST $ \r s -> f (m r s)

instance MFunctor MaybeT where
  hoist f (MaybeT m) = MaybeT (f m)

instance MFunctor (ExceptT e) where
  hoist f (ExceptT m) = ExceptT (f m)

instance MFunctor Yoneda where
  hoist f (Yoneda m) = Yoneda (f . m)

instance MFunctor Coyoneda where
  hoist f (Coyoneda k m) = Coyoneda k (f m)

-- | the terminal monad homomorphism
collapse :: m ~> Proxy
collapse _ = Proxy

instance MonadIO Proxy where
  liftIO _ = Proxy

class Commute s t where
  -- 'commute is' a monad homomorphism
  -- 
  -- At this time, @'commute'.'commute' = 'id'@ is _not_ a law.
  --
  -- @
  -- commute.commute.commute = commute
  -- @
  commute :: Monad m => s (t m) ~> t (s m)

-- * ReaderT/WriterT/StateT all commute:

instance MFunctor s => Commute s (ReaderT e) where
  commute (srema :: s (ReaderT e m) a) = 
   case transform :: Monad m :- Monad (s m) of
     Sub Dict -> ReaderT $ \e -> hoist (\rema -> runReaderT rema e) srema

instance (Monoid e, MFunctor s) => Commute (WriterT e) s where
  commute (WriterT smae :: WriterT e (s m) a) = case transform :: Monad m :- Monad (s m) of
    Sub Dict -> case transform :: Monad (WriterT e m) :- Monad (s (WriterT e m)) of
      Sub Dict -> do
         (a,e) <- hoist lift smae
         lift $ writer (a,e)

-- | isomorphism
instance Commute (ReaderT e) (StateT s) where
  commute m = StateT $ \s -> ReaderT $ \e -> runStateT (runReaderT m e) s

-- | isomorphism
instance Commute (ReaderT e) (WriterT w) where
  commute m = WriterT $ ReaderT $ \e -> runWriterT (runReaderT m e)

-- | isomorphism
instance Commute (ReaderT e) (ExceptT b) where
  commute m = ExceptT $ ReaderT $ \e -> runExceptT (runReaderT m e)

-- | isomorphism
instance Commute (StateT s) (WriterT w) where
  commute m = WriterT $ StateT $ \s -> runWriterT (runStateT m s) <&> \((a, s'), e) -> ((a, e), s')

-- | split epimorphism
instance Commute (ExceptT e) (StateT s) where
  commute (ExceptT (StateT f)) = StateT $ \s -> ExceptT $ f s <&> \(ea,s') -> fmap (,s') ea

-- | split monomorphism
instance Commute (StateT s) (ExceptT e) where
  commute m = ExceptT $ StateT $ \s -> runExceptT (runStateT m s) <&> \case
    Left e -> (Left e, s)
    Right (a, s') -> (Right a, s')

-- | split epimorphism
instance Commute (ExceptT e) (WriterT w) where
  commute (ExceptT (WriterT m)) = WriterT $ ExceptT $ m <&> \(ea,s') -> fmap (,s') ea

-- * ExceptT commutes with State

-- | Monic unless m is terminal
generalize :: Monad m => Identity ~> m
generalize (Identity a) = return a

-- | Split Epi. 'evalStateT' as a monad homomorphism
--
-- @
-- evaluate . readOnly = id
-- @
evaluate :: Functor m => StateT s m ~> ReaderT s m
evaluate = ReaderT . flip lowerStateT

-- | Split Mono. Run a reader computation using the current state as its environment
readOnly :: Functor m => ReaderT s m ~> StateT s m
readOnly m = StateT $ \s -> (,s) <$> runReaderT m s

-- | Split Epi
--
-- @
-- lowerWriterT . lift = id
-- lowerWriterT . writer = return . fst
-- @
lowerWriterT :: Functor m => WriterT e m ~> m
lowerWriterT m = fst <$> runWriterT m 

-- | Split Epi
--
-- @
-- lowerReaderT e . lift = id
-- lowerReaderT e (reader f) = return (f e)
-- lowerReaderT = flip runReaderT
-- @
lowerReaderT :: e -> ReaderT e m ~> m
lowerReaderT e m = runReaderT m e

-- | Split Epi
--
-- @
-- lowerStateT s . lift = id
-- lowerStateT = flip evalStateT
-- @
lowerStateT :: Functor m => s -> StateT s m ~> m
lowerStateT s m = fst <$> runStateT m s


-- | A Monad on the category of Monads over Hask.
--
-- 'lift' is analogous to 'return'
-- 'squash' is analogius to 'join'
-- 'hoist' is analogous to 'fmap'
-- 'embed' is analogous to ('=<<')
class (MFunctor t, MPointed t) => MMonad t where
  -- | Split Epimorphism
  --
  -- @
  -- squash . lift = id
  -- squash . hoist lift = id
  -- squash . hoist squash = squash . squash
  -- @
  --
  squash :: Monad m => t (t m) ~> t m

instance MMonad (ReaderT e) where
  squash m = ReaderT $ \e -> runReaderT (runReaderT m e) e

instance Monoid w => MMonad (WriterT w) where
  squash wwma = WriterT $ runWriterT (runWriterT wwma) <&> \((a,e),e') -> (a, e `mappend` e')
       
instance MMonad IdentityT where
  squash = runIdentityT

instance MMonad MaybeT where
  squash (MaybeT (MaybeT m)) = MaybeT (join <$> m)

instance MMonad Yoneda where
  squash = lowerYoneda

instance MMonad Coyoneda where
  squash = lowerCoyoneda

embed :: forall t m n. (MMonad t, Monad m, Monad n) => (m ~> t n) -> t m ~> t n
embed f = case transform :: Monad m :- Monad (t m) of
  Sub Dict -> case transform :: Monad n :- Monad (t n) of
    Sub Dict -> squash . hoist f

--------------------------------------------------------------------------------
-- * Tensoring transformers
--------------------------------------------------------------------------------

-- |
-- 'Tensor'/'runTensor' form a monad isomorphism
--
-- @
-- Tensor :: s (t m) ~> Tensor s t m
-- runTensor :: Tensor s t m ~> s (t m)
-- @

newtype Tensor
  (s :: (* -> *) -> * -> *)
  (t :: (* -> *) -> * -> *)
  (m :: * -> *)
  (a :: *) = Tensor { runTensor :: s (t m) a }

instance (MonadTrans s, MPointed t) => MonadTrans (Tensor s t) where
  lift (m :: m a) = case transform :: Monad m :- Monad (t m) of
    Sub Dict -> Tensor (lift (lift m))

instance (Lifting Monad s, Lifting Monad t) => Lifting Monad (Tensor s t) where
  lifting = Sub Dict

instance (MFunctor s, MFunctor t) => MFunctor (Tensor s t) where
  hoist = go where
    go :: forall m n a. (Monad m, Monad n) => (forall x. m x -> n x) -> Tensor s t m a -> Tensor s t n a
    go f (Tensor stma) = case transform :: Monad m :- Monad (t m) of
      Sub Dict -> case transform :: Monad n :- Monad (t n) of
        Sub Dict -> Tensor (hoist (hoist f) stma)

instance (Lifting Monad s, Lifting Monad t, Monad m) => Functor (Tensor s t m) where
  fmap f (Tensor m) = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (fmap f m)
  a <$ Tensor m = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (a <$ m)

instance (Lifting Monad s, Lifting Monad t, Monad m) => Applicative (Tensor s t m) where
  pure a = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (return a)
  m <*> n = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (runTensor m <*> runTensor n)
  m *> n = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (runTensor m *> runTensor n)
  m <* n = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (runTensor m <* runTensor n)

instance (Lifting Monad s, Lifting Monad t, Monad m) => Monad (Tensor s t m) where
  return = pure
  m >>= f = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (runTensor m >>= runTensor . f)
  fail s = case transform.transform :: Monad m :- Monad (s (t m)) of
    Sub Dict -> Tensor (fail s)

instance (Lifting Monad s, Lifting Monad t, Alternative (s (t m)), Monad m) => Alternative (Tensor s t m) where
  Tensor m <|> Tensor n = Tensor (m <|> n)
  empty = Tensor empty

instance (Lifting Monad s, Lifting Monad t, MonadPlus (s (t m)), Monad m) => MonadPlus (Tensor s t m) where
  mplus (Tensor m) (Tensor n) = Tensor (m `mplus` n)
  mzero = Tensor mzero

instance (Lifting Monad s, Lifting Monad t, MonadState e (s (t m)), Monad m) => MonadState e (Tensor s t m) where
  get = Tensor get
  put = Tensor . put
  state = Tensor . state

instance (Lifting Monad s, Lifting Monad t, MonadWriter e (s (t m)), Monad m) => MonadWriter e (Tensor s t m) where
  tell = Tensor . tell
  listen = Tensor . listen . runTensor
  pass = Tensor . pass . runTensor
  writer = Tensor . writer

instance (Lifting Monad s, Lifting Monad t, MonadReader e (s (t m)), Monad m) => MonadReader e (Tensor s t m) where
  ask = Tensor ask
  local f = Tensor . local f . runTensor

-- | lambda/unlambda form a monad isomorphism
lambda :: t m ~> Tensor IdentityT t m
lambda = Tensor . IdentityT

unlambda :: Tensor IdentityT t m ~> t m
unlambda = runIdentityT . runTensor

-- | rho/unrho form a monad isomorphism
rho :: (MFunctor t, Monad m) => t m ~> Tensor t IdentityT m
rho = Tensor . hoist IdentityT

unrho :: (MFunctor t, Monad m) => Tensor t IdentityT m ~> t m
unrho = hoist runIdentityT . runTensor

-- | assoc/unassoc form a monad isomorphism
assoc :: forall s t u m. (MFunctor s, MPointed t, MPointed u, Monad m) => Tensor (Tensor s t) u m ~> Tensor s (Tensor t u) m
assoc (Tensor (Tensor m)) = case transform.transform :: Monad m :- Monad (t (u m)) of
  Sub Dict -> Tensor (hoist Tensor m)

unassoc :: forall s t u m. (MFunctor s, MPointed t, MPointed u, Monad m) => Tensor s (Tensor t u) m ~> Tensor (Tensor s t) u m
unassoc (Tensor m) = case transform.transform :: Monad m :- Monad (t (u m)) of
  Sub Dict -> Tensor (Tensor (hoist runTensor m))

--------------------------------------------------------------------------------
-- * The initial monad transformer
--------------------------------------------------------------------------------

-- | IdentityT is the initial monad transformer
--
-- this is a monomorphism unless t is terminal.
generalizeT :: (MonadTrans t, Monad m) => IdentityT m ~> t m
generalizeT (IdentityT m) = lift m

--------------------------------------------------------------------------------
-- * The terminal monad transformer
--------------------------------------------------------------------------------

-- | ProxyT is the terminal monad transformer
data ProxyT m a = ProxyT

instance MonadTrans ProxyT where
  lift _ = ProxyT

instance Functor (ProxyT m) where
  fmap _ _ = ProxyT
  _ <$ _ = ProxyT

instance Applicative (ProxyT m) where
  pure _ = ProxyT
  _ <*> _ = ProxyT

instance Monad (ProxyT m) where
  return _ = ProxyT
  _ >>= _  = ProxyT
  fail _   = ProxyT

instance MonadIO (ProxyT m) where
  liftIO _ = ProxyT

-- this would ideally be an instance of MonadState s for all s, but there is a fundep!
{-
instance MonadState s (ProxyT m) where
  get = ProxyT
  put _ = ProxyT
  state _ = ProxyT
-}

instance Alternative (ProxyT m) where
  empty = ProxyT
  _ <|> _ = ProxyT

instance MonadPlus (ProxyT m) where
  mzero = ProxyT
  mplus _ _ = ProxyT

instance MonadFix (ProxyT m) where
  mfix _ = ProxyT

instance MonadZip (ProxyT m) where
  munzip _ = (ProxyT, ProxyT)
  mzipWith _ _ _ = ProxyT

instance Lifting Monad ProxyT where
  lifting = Sub Dict

instance MMonad ProxyT where
  squash _ = ProxyT

instance MFunctor ProxyT where
  hoist _ _ = ProxyT

collapseT :: m ~> ProxyT n
collapseT _ = ProxyT

--------------------------------------------------------------------------------
-- * Misc.
--------------------------------------------------------------------------------

class Trivial a
instance Trivial a

class PairOb t () => MPair (t :: * -> (* -> *) -> * -> *) where
  type PairOb t :: * -> Constraint
  -- | pair/unpair form an isomorphism
  --
  -- @
  -- pair . unpair = id
  -- unpair . pair = id
  -- @
  pair   :: (Monad m, PairOb t a, PairOb t b) => t a (t b m) ~> t (a, b) m
  unpair :: (Monad m, PairOb t a, PairOb t b) => t (a, b) m ~> t a (t b m)
  pairOb :: proxy t -> (PairOb t a, PairOb t b) :- PairOb t (a, b)
  pairT  :: PairOb t a :- MPointed (t a)
  -- | lowerUnit/lift form an isomorphism
  --
  -- @
  -- lift . lowerUnit = id
  -- lowerUnit . lift = id
  -- @
  lowerUnit :: Monad m => t () m ~> m

assocPair :: ((a,b),c) -> (a,(b,c))
assocPair ((a,b),c) = (a,(b,c))

disassocPair :: (a,(b,c)) -> ((a,b),c)
disassocPair (a,(b,c)) = ((a,b),c)

instance MPair StateT where
  type PairOb StateT = Trivial
  pair m = StateT $ \(a,b) -> assocPair <$> runStateT (runStateT m a) b
  unpair m = StateT $ \a -> StateT $ \b -> disassocPair <$> runStateT m (a,b)
  pairOb _ = Sub Dict
  pairT = Sub Dict
  lowerUnit m = evalStateT m ()
  
instance MPair WriterT where
  type PairOb WriterT = Monoid
  pair m = WriterT $ assocPair <$> runWriterT (runWriterT m)
  unpair m = WriterT $ WriterT $ disassocPair <$> runWriterT m
  pairOb _ = Sub Dict
  pairT = Sub Dict
  lowerUnit m = fst <$> runWriterT m

instance MPair ReaderT where
  type PairOb ReaderT = Trivial
  pair m = ReaderT $ \(a,b) -> runReaderT (runReaderT m a) b
  unpair m = ReaderT $ \a -> ReaderT $ \b -> runReaderT m (a, b)
  pairOb _ = Sub Dict
  pairT = Sub Dict
  lowerUnit m = runReaderT m ()

class SumOb t Void => MSum (t :: * -> (* -> *) -> * -> *) where
  type SumOb t :: * -> Constraint
  -- | sum/unsum form an isomorphism
  --
  -- @
  -- sum . unsum = id
  -- unsum . sum = id
  -- @
  sum   :: (Monad m, SumOb t a, SumOb t b) => t a (t b m) ~> t (Either a b) m
  unsum :: (Monad m, SumOb t a, SumOb t b) => t (Either a b) m ~> t a (t b m)
  sumOb :: proxy t -> (SumOb t a, SumOb t b) :- SumOb t (Either a b)
  sumT  :: SumOb t a :- MPointed (t a)
  -- | lowerCounit/lift form an isomorphism
  --
  -- @
  -- lift . lowerCounit = id
  -- lowerCounit . lift = id
  -- @
  lowerCounit :: Monad m => t Void m ~> m

instance MSum ExceptT where
  type SumOb ExceptT = Trivial
  sum m = ExceptT $ swizzle <$> runExceptT (runExceptT m) where
    swizzle (Left b) = Left (Right b)
    swizzle (Right (Left a)) = Left (Left a)
    swizzle (Right (Right c)) = Right c
  unsum m = ExceptT $ ExceptT $ unswizzle <$> runExceptT m where
    unswizzle (Left (Right b)) = Left b
    unswizzle (Left (Left a)) = Right (Left a)
    unswizzle (Right c) = Right (Right c)
  sumOb _ = Sub Dict
  sumT = Sub Dict
  lowerCounit m = either absurd id <$> runExceptT m