A reimplementation of monad-control with an alternate API

Control.hs

{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}

module Control.Monad.Trans.Control where

import Control.Monad (void)
import Control.Monad.Base (MonadBase, liftBase)
import Control.Monad.Except (ExceptT(..))
import Control.Monad.Reader (ReaderT(..), ask)
import Control.Monad.State (StateT(..), get)
import Control.Monad.Trans (MonadTrans, lift)
import Control.Monad.Writer (WriterT(..))

class MonadTrans t => MonadTransControl t where
  type CtxT t :: *
  type StT t a :: *

  -- instances should satisfy the following law:
  --   restoreT . runT x =<< ctxT ≡ x

  ctxT :: Monad m => t m (CtxT t)
  runT :: Monad m => t m a -> CtxT t -> m (StT t a)
  restoreT :: Monad m => m (StT t a) -> t m a

liftWith :: (MonadTransControl t, Monad m, Monad (t m)) => ((forall n b. Monad n => t n b -> n (StT t b)) -> m a) -> t m a
liftWith f = do
  ctx <- ctxT
  lift $ f (flip runT ctx)

--------------------------------------------------------------------------------

class MonadBase b m => MonadBaseControl b m | m -> b where
  type CtxM m :: *
  type StM m a :: *

  ctxM :: m (CtxM m)
  runM :: m a -> CtxM m -> b (StM m a)
  restoreM :: b (StM m a) -> m a

  default ctxM :: (MonadTransControl t, MonadBaseControl b m', m ~ t m', CtxM m ~ ComposeCtx t m') => m (CtxM m)
  ctxM = (,) <$> ctxT <*> lift ctxM

  default runM :: (MonadTransControl t, MonadBaseControl b m', m ~ t m', CtxM m ~ ComposeCtx t m', StM m a ~ ComposeSt t m' a) => m a -> CtxM m -> b (StM m a)
  runM m (ctx, ctx') = runM (runT m ctx) ctx'

  default restoreM :: (MonadTransControl t, MonadBaseControl b m', m ~ t m', StM m a ~ ComposeSt t m' a) => b (StM m a) -> m a
  restoreM = restoreT . restoreM

type ComposeCtx t m = (CtxT t, CtxM m)
type ComposeSt t m a = StM m (StT t a)

liftBaseWith :: MonadBaseControl b m => ((forall c. m c -> b (StM m c)) -> b a) -> m a
liftBaseWith f = do
  ctx <- ctxM
  liftBase $ f (flip runM ctx)

liftBaseOp :: MonadBaseControl b m => ((a -> b (StM m c)) -> b (StM m d)) -> (a -> m c) -> m d
liftBaseOp op f = do
  ctx <- ctxM
  restoreM (op (\x -> runM (f x) ctx))

liftBaseOp_ :: MonadBaseControl b m => (b (StM m a) -> b (StM m c)) -> m a -> m c
liftBaseOp_ op m = do
  ctx <- ctxM
  restoreM (op (runM m ctx))

liftBaseDiscard :: MonadBaseControl b m => (b () -> b a) -> m () -> m a
liftBaseDiscard op m = do
  ctx <- ctxM
  liftBase (op (void $ runM m ctx))

liftBaseOpDiscard :: MonadBaseControl b m => ((a -> b ()) -> b c) -> (a -> m ()) -> m c
liftBaseOpDiscard op f = do
  ctx <- ctxM
  liftBase (op (\x -> void $ runM (f x) ctx))

--------------------------------------------------------------------------------

instance MonadTransControl (ReaderT r) where
  type CtxT (ReaderT r) = r
  type StT (ReaderT r) a = a

  ctxT = ask
  runT (ReaderT f) = f
  restoreT = ReaderT . const

instance Monoid w => MonadTransControl (WriterT w) where
  type CtxT (WriterT w) = ()
  type StT (WriterT w) a = (a, w)

  ctxT = return ()
  runT (WriterT m) = const m
  restoreT = WriterT

instance MonadTransControl (StateT s) where
  type CtxT (StateT s) = s
  type StT (StateT s) a = (a, s)

  ctxT = get
  runT (StateT f) = f
  restoreT = StateT . const

instance MonadTransControl (ExceptT e) where
  type CtxT (ExceptT e) = ()
  type StT (ExceptT e) a = Either e a

  ctxT = return ()
  runT (ExceptT m) = const m
  restoreT = ExceptT

--------------------------------------------------------------------------------

instance MonadBaseControl b m => MonadBaseControl b (ReaderT r m) where
  type CtxM (ReaderT r m) = ComposeCtx (ReaderT r) m
  type StM (ReaderT r m) a = ComposeSt (ReaderT r) m a

instance (MonadBaseControl b m, Monoid w) => MonadBaseControl b (WriterT w m) where
  type CtxM (WriterT w m) = ComposeCtx (WriterT w) m
  type StM (WriterT w m) a = ComposeSt (WriterT w) m a

instance MonadBaseControl b m => MonadBaseControl b (StateT s m) where
  type CtxM (StateT s m) = ComposeCtx (StateT s) m
  type StM (StateT s m) a = ComposeSt (StateT s) m a

instance MonadBaseControl b m => MonadBaseControl b (ExceptT e m) where
  type CtxM (ExceptT e m) = ComposeCtx (ExceptT e) m
  type StM (ExceptT e m) a = ComposeSt (ExceptT e) m a