Unexpected output of callCC in Haskell

Question

I was trying to build a very simple example of a coroutine in Haskell using callCC. My understanding of how callCC :: ((a -> m b) -> m a) -> m a works is that by writing

example :: Cont Int Int
example = do
   callCC $ \k -> ...
   f

I essentially bind the continutation of that callCC, i.e. f to k. Then, if k is used inside of callCC's body, e.g. callCC $ \k -> ... k 42 ... it will result in the continuation being called with this value, i.e., f 42, otherwise, if k is unused, callCC $ \k -> ... return 41, it will result in f applied to 41 instead. By storing this k in some data structure I gain the ability of resuming the function (example) right after the callCC (because f is bound to k).

The code I have for a coroutine is as follows

{-#LANGUAGE GeneralisedNewtypeDeriving #-}

module Cont where

import Control.Monad.Cont
import Control.Monad.IO.Class(liftIO)
import Control.Monad.State(put, get, runStateT, evalStateT, StateT, MonadState, MonadIO, execStateT)

newtype CoroutineT r m a = CoroutineT {runCoroutineT :: ContT r (StateT [CoroutineT r m ()] m) a}
    deriving (Functor,Applicative,Monad,MonadCont,MonadIO, MonadState [CoroutineT r m ()])


schedule :: CoroutineT () IO ()
schedule = do
    cs <- get
    case cs of
        (c:cs) -> do
            put cs
            c
            schedule
        [] -> return ()

yield :: CoroutineT () IO ()
yield = callCC $ \k1 -> do
            cs <- get
            put (cs ++ [k1 ()])
            schedule

test1 :: CoroutineT () IO ()
test1 = do
    liftIO $ print "HelloTest1"
    yield
    liftIO $ print "HelloTest1_2"

test2 :: CoroutineT () IO ()
test2 = do
    liftIO $ print "HelloTest2"
    yield
    liftIO $ print "HelloTest2_2"

test_cont :: IO ()
test_cont = do 
    runStateT (evalContT (runCoroutineT schedule)) [test1,test2]
    return ()

where yield stores the continuation in the state which is then resumed by schedule.

I would expect each continuation to be executed twice --- once in schedule after 'invoking' c and then after each yeild gets evaluated. And therefore the expected output should be

"HelloTest1"
"HelloTest2"
"HelloTest1_2"
"HelloTest2_2"
"HelloTest2_2"
"HelloTest1_2"

What happens in reality though is

"HelloTest1"
"HelloTest2"
"HelloTest1_2"
"HelloTest2_2"
"HelloTest1_2"

What might be wrong with my understanding and/or the code?

UPD:

As correctly noted by @Li-yaoXia I was missing the fact that the continuation after invoking k bound inside of callCC is ignored. It can be seen from the definition of a callCC

callCC :: ((a -> Cont r b) -> Cont r a) -> Cont r a
callCC f = Cont $ \k1 -> let a2b = \x -> Cont $ \_ -> k1 x in
                         runCont $ (f a2b) k1

when I invoke the callCC as callCC $ \k -> ... the k is bound to a2b and once I invoke this k outside the callCC it becomes \x -> Cont $ \_ -> k1 x where k1 is the continuation after the callCC and _ is the continuation after the invocation of k which is ignored.

Answer 1

You're probably missing the fact that when a continuation captured with callCC is called (c in schedule), the current continuation at the call site (schedule and whatever follows) is discarded.

PROGRAM            | STATE
---------------------------
schedule             [test1,test2]
---------------------------
test1                [test2]
schedule
--------------------------- Print "HelloTest1"
yield                [test2]
print "HelloTest1_2"
schedule
---------------------------  -- let K1 = do print "HelloTest1_2"; schedule
schedule             [test2, K1]
print "HelloTest1_2"
schedule
---------------------------
test2                [K1]
print "HelloTest1_2"
schedule
--------------------------- Print "HelloTest2"
yield                [K1]
print "HelloTest2_2"
schedule
print "HelloTest1_2"
schedule
---------------------------  -- let K2 = print "HelloTest2_2"; schedule; print "HelloTest1_2"; schedule
schedule             [K1, K2]
print "HelloTest2_2"
schedule
print "HelloTest1_2"
schedule
--------------------------- -- call K1: the current continuation is discarded, the whole program becomes K1 (Note the definition of K1 above is a lie since it doesn't include this detail.)
print "HelloTest1_2" [K2]
schedule
--------------------------- Print "HelloTest1_2"
schedule             [K2]
--------------------------- -- call K2: the current continuation is discarded, the whole program becomes K2
print "HelloTest2_2" []
schedule
print "HelloTest1_2"
schedule
--------------------------- Print "HelloTest2_2" and "HelloTest1_2"