import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Class
import Control.Monad.Identity

newtype TraceT t m a = TraceT { runTraceT :: (t -> m ()) -> m a }

instance (Functor m) => Functor (TraceT t m) where
  fmap f ma = TraceT $ \tr -> f <$> runTraceT ma tr

instance (Applicative m) => Applicative (TraceT t m) where
  pure = liftTraceT . pure
  mf <*> ma = TraceT $ \tr -> runTraceT mf tr <*> runTraceT ma tr

instance (Monad m) => Monad (TraceT t m) where
  ma >>= f = TraceT $ \tr -> do
    a <- runTraceT ma tr
    runTraceT (f a) tr

instance MonadTrans (TraceT t) where
  lift = liftTraceT

liftTraceT :: m a -> TraceT t m a
liftTraceT = TraceT . const

trace :: t -> TraceT t m ()
trace t = TraceT ($t)

mute :: Applicative m => TraceT t m a -> TraceT t m a
mute = alter (const (const (pure ())))

annotate :: Monad m => (t -> m ()) -> TraceT t m a -> TraceT t m a
annotate tr2 = alter (\tr1 -> \t -> tr1 t >> tr2 t)

alter :: ((s -> m ()) -> (t -> m ())) -> TraceT t m a -> TraceT s m a
alter f t = TraceT $ \tr -> runTraceT t (f tr)

mapTrace :: (s -> t) -> TraceT s m a -> TraceT t m a
mapTrace f t = TraceT $ \tr -> runTraceT t (tr . f)