Using MMonad and MonadTransControl to define (with :: (a -> Foo m x) -> Foo m x) in terms of (with :: (a -> m x) -> Foo m x)

With.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, TypeFamilies, TupleSections #-}

import Control.Exception
import Control.Applicative
import Control.Monad.State.Strict
import Control.Monad.Trans.Class
import Control.Monad.Trans.Either
import Control.Monad.Error
import Control.Monad.Trans.Control
import Control.Monad.Morph
import Control.Error.Util
import Data.Either.Combinators
import Data.Monoid

import Debug.Trace

type Info = [Char]
data Err = Whatever | IOE Info (IOException)
  deriving Show

instance Error Err where
  noMsg = Whatever

newtype Foo m a = Foo { foo :: ErrorT Err (StateT Info m) a }
  deriving ( Functor, Applicative, Monad
           , MonadError Err, MonadIO, (MonadState Info)
           ) --, MonadTrans, MFunctor, MMonad)

instance MonadTrans Foo where
  lift = lift

instance MFunctor Foo where
  hoist = hoist

-- this is needed for "squash"
instance MMonad Foo where
  embed f m = Foo $ ErrorT $ StateT $ \s -> do
    (a, s1) <- runStateT (runErrorT $ foo $ f (runStateT (runErrorT $ foo $ m) s)) s
    case a of Left e -> return (Left e, s1)
              Right (Left e, s2) -> return (Left e, mappend s2 s1)
              Right (Right x, s2) -> return (Right x, mappend s2 s1)

instance MonadTransControl Foo where
  data StT Foo a = StFoo { unStFoo :: (Either Err a, Info) }
  liftWith f = Foo $ ErrorT $ StateT $ \s ->
    liftM (, s)                                                            -- rewrap state
          (liftM return                                                    -- rewrap error
                 (f $ \t -> liftM StFoo
                                  (runStateT (runErrorT $ foo t) s))) -- unwrap error and state
  restoreT = Foo . ErrorT . StateT . const . liftM unStFoo
  {-# INLINE liftWith #-}
  {-# INLINE restoreT #-}

unwrap = fmap StFoo . flip runStateT [] . runErrorT . foo

with :: (Int -> IO a) -> Foo IO a
with act = trace "with" $ liftIO $ act 0

with' :: (Int -> Foo IO a) -> Foo IO a
with' act = trace "with'" $ squash $ restoreT $ do
  with (unwrap . act)

fizz :: Foo IO ()
fizz = do
  x <- get
  put ('a':x)
  with (\_ -> print "put a")

buzz = do
  liftIO $ print "buzz"
  fizz
  with' (\_ -> do liftIO $ print "getting the x"
                  x <- get
                  put ('b':x))
  fizz

main = do
  s <- execStateT (runErrorT (foo buzz)) []
  putStrLn (show s)