Using monad transformers and continuations to make a minimal interpreter for procedure early return

Question

For the purpose of learning, I'm doing an interpreter for a minimal language with subprocedure calls and returns.

data P = Px Int | Ps [P] | Pc P | Pr

The meaning is: Px x instruction x, Ps xs sequence of instructions, Pc x call subprocedure x, Pr early return. The configuration during a computation is just an Int and the instruction Px increments it (just a way to visualize this minimal example, I'm actually applying this thing to a bigger language). So for example run 0 $ Ps [Px 1, Px 2] = [1,3] is a complete trace of execution starting from the configuration 0.

I realized that to handle the early return I need continuations, so I did the following

runm :: Int -> P -> ([Int] -> Cont [Int] [Int]) -> Cont [Int] [Int]
runm c p k = case p of
  Px x -> return [c+x]
  Ps [] -> return []
  Ps (x:xs) -> do
    s <- runm c x k
    let c2 = if not (null s) then last s else c
        k' r = k $ s ++ r
    ss <- runm c2 (Ps xs) k'
    return $ s ++ ss
  Pc x -> callCC $ runm c x
  Pr -> k []

It seems that it works correctly, for example:

> evalCont $ runm 0 (Ps [Px 1, Pc $ Ps [Px 2, Pr, Px 100], Px 3]) undefined
[1,3,6]

Executes "1", "2" and "3" but correctly skips "100" as it comes after a return.

It bothers me a bit in this situation that I have to somehow manage the escaper continuation in the Ps case, that hasn't so much to do with the return. I therefore thought that a writer monad might be of help.

The idea is that the writer monad would take care of "appending" successive configurations while the cont handles the "backward jump" of return.

Being not very familiar with continuations and monad transformers I haven't had much success anyway. I can't even really understand in which order should I build the monad stack.

For example:

runwc :: Int -> P -> ([Int] -> ContT [Int] (Writer [Int]) [Int]) -> ContT [Int] (Writer [Int]) [Int]
runwc c p k = case p of
  Px x -> (lift $ tell [c+x]) >> return []
  Ps [] -> return []
  Ps (x:xs) -> do
    (_,s) <- lift $ listen $ evalContT $ runwc c x k
    let c2 = if not (null s) then last s else c
    runwc c2 (Ps xs) k
  Pc x -> callCC $ runwc c x
  Pr -> k []

This doesn't really return:

> execWriter $ evalContT $ runwc 0 (Ps [Px 1, Pc $ Ps [Px 2, Pr, Px 100], Px 3]) undefined
[1,3,103,106]

I think I didn't understand exactly why.

The opposite order was seeming to me more promising, but this doesn't work either:

runcw :: Int -> P -> (() -> WriterT [Int] (Cont [Int]) ()) -> WriterT [Int] (Cont [Int]) ()
runcw c p k = case p of
  Px x -> tell [c+x]
  Ps [] -> return ()
  Ps (x:xs) -> do
    (_,s) <- listen $ runcw c x k
    let c2 = if not (null s) then last s else c
    runcw c2 (Ps xs) k
  Pc x -> WriterT $ callCC $ \j -> runWriterT $ do
    let k' = \_ -> WriterT $ j ((), [])
    runcw c x k'
  Pr -> k ()

It returns too much:

> evalCont $ execWriterT $ runcw 0 (Ps [Px 1, Pc $ Ps [Px 2, Pr, Px 100], Px 3]) undefined
[1,4]

The difference is that here I think I understand it better: the escaper function j ought to be invoked with the collected history from the subsequent call to runwc c x k' but it isn't getting it. In j ((), []) everything is discarded and callCC returns an empty result for the Pc case.

My idea was that the writer monad would "independently" collect the trace so it would be present even when jumping through the escaper continuations, but it seems it can't work like this, since there doesn't seem to be a way for j to get the "past" (I tried some recursion from the subsequent runcw call but it looped).

To clarify what I expected to be able to do, I can do it using the more powerful state monad instead:

runcs :: Int -> P -> (() -> StateT [Int] (Cont [Int]) ()) -> StateT [Int] (Cont [Int]) ()
runcs c p k = case p of
  Px x -> modify (++ [c+x])
  Ps [] -> return ()
  Ps (x:xs) -> do
    s <- runcs c x k >> get
    let c2 = if not (null s) then last s else c
    runcs c2 (Ps xs) k
  Pc x -> StateT $ \s -> callCC $ \j -> flip runStateT s $ do
    let k' = \_ -> StateT $ \s' -> j ((), s')
    runcs c x k'
  Pr -> k ()

This works correctly

> evalCont $ execStateT (runcs 0 (Ps [Px 1, Pc $ Ps [Px 2, Pr, Px 100], Px 3]) undefined) []
[1,3,6]

This last solution frees me from handling the escaper continuation in the Ps sequence case and it's able to state-get what happened up to the return to throw it to j.

The problem is that this is evidently too powerful: at any point the complete execution trace can be accessed and manipulated through the "global" state.

Is it possible to obtain the advantages of the State solution by using only a Writer so that each step of the interpreter can do only the small part of appending its result?

My impression is that the correct signature will actually be the ContT r (Writer w) a, even if I got more progress with the WriterT w (ContT r) a. The first corresponds to (a -> (r,w)) -> (r,w), while the second to ((a,w) -> r) -> r and it seems to me the this last carries too much symmetry between a and w, but about here my head started the explosion procedure and so I'm asking, as this is actually the first time that I do something meaningful with continuations besides trivial tests.

Answer 1

I ended up with this, which I find to be satisfactorily simpler than the original code, having moved all the details under the monadic stack. (Well, if we accept "hiding the complexity" as "simplifying", at least. :-P )

import Control.Monad.Cont
import Control.Monad.Trans.RWS

data P = Px Int | Ps [P] | Pc P | Pr

-- A shorthand
type T = ContT () (RWS () [Int] Int) ()

runwc :: P -> (() -> T) -> T
runwc p k = case p of
  Px x -> lift $ do
     c <- get
     tell [c+x]
     put (c+x)
  Ps xs -> mapM_ (flip runwc k) xs
  -- Equivalent to:
  -- Ps [] -> return ()
  -- Ps (x:xs) -> do
  --   runwc x k
  --   runwc (Ps xs) k
  Pc x -> callCC $ runwc x
  Pr -> k ()

test :: [Int]
test = trace
   where
   (_result, _counter, trace) = runRWS action () 0
   action = runContT (runwc (Ps [Px 1, Pc $ Ps [Px 2, Pr, Px 100], Px 3]) return)
            (const $ return ())

The output is

> test
[1,3,6]

as intended.

The main monadic type T needs some comment:

type T = ContT () (RWS () [Int] Int) ()
            -- 1       2  3     4    5

Here:

1. () is the type of the final result, once the continuation has been applied. I chose () since the trace is in the "writer" monad, but could be something else.
2. () is the type of the read-only state (the "reader" part of RWS). This is trivial since we don't need that.
3. [Int] the type of the trace (the "writer" part of RWS).
4. Int the type of the counter (the "state" part of RWS).
5. () the type of the result of the monadic action, the one which will be passed to the continuation (could be something else).

The rest of the code should be more or less clear. Px gets the state, increments it, and logs it. Ps is trivial: we call runwc p k for each p in the block. Pc sets the current continuation. Pr calls the set continuation.