{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE FlexibleInstances #-}

module Core.Control.Monad.Trans
  ( MonadHoist (..)
  , MonadTransBridge (..)
  , lift2
  , rehoist
  , mixinEffect
  ) where

import Data.Functor.Identity
import Control.Monad.Trans.Class
import Control.Monad.Trans.Writer.Strict

-- | Can transform the stacked version by transforming the unstacked
-- version, without knowing the unstacked version's type.
class (MonadTrans t) => MonadHoist t where
  mapInner :: (forall b. w1 b -> w2 b) -> t w1 a -> t w2 a

-- | Bridges a monad and its transformer variant.
class (Monad u, MonadHoist t) => MonadTransBridge u t where
  hoist :: (Monad w) => u a -> t w a
  bindStackOuter :: (Monad w) => (u a -> t w b) -> t w a -> t w b

instance (Monoid r) => MonadHoist (WriterT r) where
  mapInner = mapWriterT

instance (Monoid r) => MonadTransBridge (WriterT r Identity) (WriterT r) where
  hoist = mapWriterT $ pure . runIdentity
  bindStackOuter f = mapWriterT $ (runWriterT . f . writer =<<)

-- | Wraps an effect in 2 transformers.
lift2 :: (MonadTrans t1, MonadTrans t2, Monad u, Monad (t2 u))
      => u a
      -> t1 (t2 u) a
lift2 = lift . lift

-- | Hoists the hoisted monad.
rehoist :: (MonadHoist t, MonadTransBridge wu wt, Monad ww)
        => t wu a
        -> t (wt ww) a
rehoist = mapInner hoist

-- | Apply the effect, then return the original value.
mixinEffect :: (MonadTransBridge w wt, Monad we, Monad (wt we))
            => (w a -> we ())
            -> w a
            -> wt we a
mixinEffect eff x = lift (eff x) >> hoist x