The Pause monad

Monads can do many amazing, crazy things. They can create variables which hold a superposition of values. They can allow you to access data from the future before you compute it. They can allow you to write destructive updates, but not really. And then the continuation monad allows you to break people's minds! Ususally your own. ;-)

But here's a challenge: Can you make a monad which can be paused?

data Pause s x
instance Monad (Pause s)
mutate :: (s -> s) -> Pause s ()
yield :: Pause s ()
step :: s -> Pause s () -> (s, Maybe (Pause s ()))

The Pause monad is a kind of state monad (hence mutate, with the obvious semantics). Normally a monad like this has some sort of "run" function, which runs the computation and hands you back the final state. But Pause is different: It provides a step function, which runs the computation until it calls the magical yield function. Here the computation is paused, returning to the caller enough information to resume the computation later.

For extra awesomness: Allow the caller to modify the state between step calls. (The type signatures above ought to allow this, for example.)

---

Use case: It's often easy to write code that does something complex, but a total PITA to transform it to also output the intermediate states in its operation. If you want the user to be able to change something mid-way through execution, things get complex really fast.

Implementation ideas:

• Obviously it can be done with threads, locks and IO. But can we do better? ;-)

• Something insane with a continuation monad?

• Maybe some kind of writer monad, where yield just logs the current state, and then we can "pretend" to step it by iterating over the states in the log. (Obviously this precludes altering the state between steps, since we're not really "pausing" anything now.)

Note: that you provided yourself no direct access to the current state s in this monad.

Pause s is just a free monad over the mutate and yield operations. Implemented directly you get:

data Pause s a
  = Return a
  | Mutate (s -> s) (Pause s a)
  | Yield (Pause s a)

instance Monad (Pause s) where
  return = Return
  Return a   >>= k = k a
  Mutate f p >>= k = Mutate f (p >>= k)
  Yield p    >>= k = Yield (p >>= k)

with a couple of smart constructors to give you the desired API:

mutate :: (s -> s) -> Pause s ()
mutate f = Mutate f (return ())

yield :: Pause s ()
yield = Yield (return ())

and the step function to drive it

step :: s -> Pause s () -> (s, Maybe (Pause s ()))
step s (Mutate f k) = step (f s) k
step s (Return ()) = (s, Nothing)
step s (Yield k) = (s, Just k)

You could also define this directly using

data Free f a = Pure a | Free (f (Free f a))

(from my 'free' package) with

data Op s a = Mutate (s -> s) a | Yield a

then we already have a monad for Pause

type Pause s = Free (Op s)

and just need to define the smart constructors and stepper.

Making it faster.

Now, these implementations are easy to reason about, but they don't have the fastest operational model. In particular, left associated uses of (>>=) yield asymptotically slower code.

To get around that you can apply the Codensity monad to your existing free monad, or just use the 'Church free' monad directly, both of which I describe in depth on my blog.

http://comonad.com/reader/2011/free-monads-for-less/

http://comonad.com/reader/2011/free-monads-for-less-2/

http://comonad.com/reader/2011/free-monads-for-less-3/

The result of applying the Church encoded version of the Free monad is that you get an easy to reason about model for the data type, and you still get a fast evaluation model.

Sure; you just let any computation either finish with a result, or suspend itself, giving an action to be used on resume, along with the state at the time:

data Pause s a = Pause { runPause :: s -> (PauseResult s a, s) }

data PauseResult s a
    = Done a
    | Suspend (Pause s a)

instance Monad (Pause s) where
    return a = Pause (\s -> (Done a, s))
    m >>= k = Pause $ \s ->
        case runPause m s of
            (Done a, s') -> runPause (k a) s'
            (Suspend m', s') -> (Suspend (m' >>= k), s')

get :: Pause s s
get = Pause (\s -> (Done s, s))

put :: s -> Pause s ()
put s = Pause (\_ -> (Done (), s))

yield :: Pause s ()
yield = Pause (\s -> (Suspend (return ()), s))

step :: Pause s () -> s -> (Maybe (Pause s ()), s)
step m s =
    case runPause m s of
        (Done _, s') -> (Nothing, s')
        (Suspend m', s') -> (Just m', s')

The Monad instance just sequences things in the normal way, passing the final result to the k continuation, or adding the rest of the computation to be done on suspension.

Here's how I'd go about it, using free monads. Er, um, what are they? They're trees with actions at the nodes and values at the leaves, with >>= acting like tree grafting.

data f :^* x
  = Ret x
  | Do (f (f :^* x))

It's not unusual to write F^*X for such a thing in the mathematics, hence my cranky infix type name. To make an instance, you just need f to be something you can map over: any Functor will do.

instance Functor f => Monad ((:^*) f) where
  return = Ret
  Ret x  >>= k  = k x
  Do ffx >>= k  = Do (fmap (>>= k) ffx)

That's just "apply k at all the leaves and graft in the resulting trees". These can trees represent strategies for interactive computation: the whole tree covers every possible interaction with the environment, and the environment chooses which path in the tree to follow. Why are they free? They're just trees, with no interesting equational theory on them, saying which strategies are equivalent to which other strategies.

Now let's have a kit for making Functors which correspond to a bunch of commands we might want to be able to do. This thing

data (:>>:) s t x = s :? (t -> x)

instance Functor (s :>>: t) where
  fmap k (s :? f) = s :? (k . f)

captures the idea of getting a value in x after one command with input type s and output type t. To do that, you need to choose an input in s and explain how to continue to the value in x given the command's output in t. To map a function across such a thing, you tack it onto the continuation. So far, standard equipment. For our problem, we may now define two functors:

type Modify s  = (s -> s) :>>: ()
type Yield     = () :>>: ()

It's like I've just written down the value types for the commands we want to be able to do!

Now let's make sure we can offer a choice between those commands. We can show that a choice between functors yields a functor. More standard equipment.

data (:+:) f g x = L (f x) | R (g x)

instance (Functor f, Functor g) => Functor (f :+: g) where
  fmap k (L fx) = L (fmap k fx)
  fmap k (R gx) = R (fmap k gx)

So, Modify s :+: Yield represents the choice between modifying and yielding. Any signature of simple commands (communicating with the world in terms of values rather than manipulating computations) can be turned into a functor this way. It's a bother that I have to do it by hand!

That gives me your monad: the free monad over the signature of modify and yield.

type Pause s = (:^*) (Modify s :+: Yield)

I can define the modify and yield commands as one-do-then-return. Apart from negotiating the dummy input for yield, that's just mechanical.

mutate :: (s -> s) -> Pause s ()
mutate f = Do (L (f :? Ret))

yield :: Pause s ()
yield = Do (R (() :? Ret))

The step function then gives a meaning to the strategy trees. It's a control operator, constructing one computation (maybe) from another.

step :: s -> Pause s () -> (s, Maybe (Pause s ()))
step s (Ret ())            = (s, Nothing)
step s (Do (L (f  :? k)))  = step (f s) (k ())
step s (Do (R (() :? k)))  = (s, Just (k ()))

The step function runs the strategy until either it finishes with a Ret, or it yields, mutating the state as it goes.

The general method goes like this: separate the commands (interacting in terms of values) from the control operators (manipulating computations); build the free monad of "strategy trees" over the signature of commands (cranking the handle); implement the control operators by recursion over the strategy trees.

Doesn't match your type signatures exactly, but certainly simple:

{-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, UndecidableInstances #-}
import Control.Monad.State

newtype ContinuableT m a = Continuable { runContinuable :: m (Either a (ContinuableT m a)) }
instance Monad m => Monad (ContinuableT m) where
    return = Continuable . return . Left
    Continuable m >>= f = Continuable $ do
        v <- m
        case v of
            Left  a -> runContinuable (f a)
            Right b -> return (Right (b >>= f))

instance MonadTrans ContinuableT where
    lift m = Continuable (liftM Left m)

instance MonadState s m => MonadState s (ContinuableT m) where
    get = lift get
    put = lift . put

yield :: Monad m => ContinuableT m a -> ContinuableT m a
yield = Continuable . return . Right

step :: ContinuableT (State s) a -> s -> (Either a (ContinuableT (State s) a), s)
step = runState . runContinuable

-- mutate unnecessary, just use modify

{-# LANGUAGE TupleSections #-}
newtype Pause s x = Pause (s -> (s, Either x (Pause s x)))

instance Monad (Pause s) where
  return x = Pause (, Left x)

  Pause k >>= f = Pause $ \s -> let (s', v) = k s in
                                case v of
                                  Left x -> step (f x) s'
                                  Right x -> (s', Right (x >>= f))

mutate :: (s -> s) -> Pause s ()
mutate f = Pause (\s -> (f s, Left ()))

yield :: Pause s ()
yield = Pause (, Right (return ()))

step :: Pause s x -> s -> (s, Either x (Pause s x))
step (Pause x) = x

That's how I would wrote this. I gave step a bit more general definition, it could be as well named runPause. In fact thinking about type of step lead me to definition of Pause.

In the monad-coroutine package you will find a general monad transformer. The Pause s monad is the same as Coroutine (State s) Id. You can combine coroutines with other monads.

Related: the Prompt monad in http://themonadreader.files.wordpress.com/2010/01/issue15.pdf

Note: This answer is available as a literate Haskell file at Gist.

I quite enjoyed this exercise. I tried to do it without looking at the answers, and it was worth it. It took me considerable time, but the result is surprisingly close to two of the other answers, as well as to monad-coroutine library. So I guess this is somewhat natural solution to this problem. Without this exercise, I wouldn't understand how monad-coroutine really works.

To add some additional value, I'll explain the steps that eventually led me to the solution.

Recognizing the state monad

Since we're dealing with states, it we look for patterns that can be effectively described by the state monad. In particular, s - s is isomorphic to s -> (s, ()), so it could be replaced by State s (). And function of type s -> x -> (s, y) can be flipped to x -> (s -> (s, y)), which is actually x -> State s y. This leads us to updated signatures

mutate :: State s () -    Pause s ()
step   :: Pause s () -    State s (Maybe (Pause s ()))

Generalization

Our Pause monad is currently parametrized by the state. However, now we see that we don't really need the state for anything, nor we use any specifics of the state monad. So we could try to make a more general solution that is parametrized by any monad:

mutate :: (Monad m) =    m () -> Pause m ()
yield  :: (Monad m) =    Pause m ()
step   :: (Monad m) =    Pause m () -> m (Maybe (Pause m ()))

Also, we could try to make mutate and step more general by allowing any kind of value, not just (). And by realizing that Maybe a is isomorphic to Either a () we can finally generalize our signatures to

mutate :: (Monad m) =    m a -> Pause m a
yield  :: (Monad m) =    Pause m ()
step   :: (Monad m) =    Pause m a -> m (Either (Pause m a) a)

so that step returns the intermediate value of the computation.

Monad transformer

Now, we see that we're actually trying to make a monad from a monad - add some additional functionality. This is what is usually called a monad transformer. Moreover, mutate's signature is exactly the same as lift from MonadTrans. Most likely, we're on the right track.

The final monad

The step function seems to be the most important part of our monad, it defines just what we need. Perhaps, this could be the new data structure? Let's try:

import Control.Monad
import Control.Monad.Cont
import Control.Monad.State
import Control.Monad.Trans

data Pause m a
    = Pause { step :: m (Either (Pause m a) a) }

If the Either part is Right, it's just a monadic value, without any suspensions. This leads us how to implement the easist thing - the lift function from MonadTrans:

instance MonadTrans Pause where
    lift k = Pause (liftM Right k)

and mutate is simply a specialization:

mutate :: (Monad m) => m () -> Pause m ()
mutate = lift

If the Either part is Left, it represents the continued computation after a suspension. So let's create a function for that:

suspend :: (Monad m) => Pause m a -> Pause m a
suspend = Pause . return . Left

Now yielding a computation is simple, we just suspend with an empty computation:

yield :: (Monad m) => Pause m ()
yield = suspend (return ())

Still, we're missing the most important part. The Monad instance. Let's fix it. Implementing return is simple, we just lift the inner monad. Implementing >>= is a bit trickier. If the original Pause value was only a simple value (Right y), then we just wrap f y as the result. If it is a paused computation that can be continued (Left p), we recursively descend into it.

instance (Monad m) => Monad (Pause m) where
    return x = lift (return x) -- Pause (return (Right x))
    (Pause s) >>= f
        = Pause $ s >>= \x -> case x of
            Right y     -> step (f y)
            Left p      -> return (Left (p >>= f))

Testing

Let's try to make some model function that uses and updates state, yielding while inside the computation:

test1 :: Int -> Pause (State Int) Int
test1 y = do
            x <- lift get
            lift $ put (x * 2)
            yield
            return (y + x)

And a helper function that debugs the monad - prints its intermediate steps to the console:

debug :: Show s => s -> Pause (State s) a -> IO (s, a)
debug s p = case runState (step p) s of
                (Left next, s')     ->  print s' >> debug s' next
                (Right r, s')       ->  return (s', r)    

main :: IO ()
main = do
    debug 1000 (test1 1 >>= test1 >>= test1) >>= print

The result is

2000
4000
8000
(8000,7001)

as expected.

Coroutines and monad-coroutine

What we have implemented is a quite general monadic solution that implements Coroutines. Perhaps not surprisingly, someone had the idea before :-), and created the monad-coroutine package. Less surprisingly, it's quite similar to what we created.

The package generalizes the idea even further. The continuing computation is stored inside an arbitrary functor. This allows suspend many variations how to work with suspended computations. For example, to pass a value to the caller of resume (which we called step), or to wait for a value to be provided to continue, etc.