Reading "shift/reset プログラミング入門"

shift-reset-2.7.hs

import Control.Monad.Cont

data Tree = Empty
          | Node Tree Int Tree

data Result = Done
            | Next Int (() -> Cont Result Result)

yield :: Int -> Cont Result ()
yield n = shiftT $ \k -> return (Next n k)

walk :: Tree -> Cont Result ()
walk Empty = return ()
walk (Node t1 n t2) = do
    walk t1
    yield n
    walk t2

start :: Tree -> Cont Result Result
start t = resetT $ walk t >> return Done

printNodes :: Tree -> IO ()
printNodes = loop . start
    where
    loop c =
        case runCont c id of 
            Done -> return ()
            Next n k -> print n >> loop (k ())

addTree :: Tree -> Int
addTree = loop . start
    where
    loop c =
        case runCont c id of
            Done -> 0
            Next n k -> n + loop (k ())

sameFringe :: Tree -> Tree -> Bool
sameFringe t1 t2 = loop (start t1) (start t2)
    where
    loop c1 c2 = 
        case (runCont c1 id, runCont c2 id) of
            (Done, Done) -> True
            (Done, _)    -> False
            (_, Done)    -> False
            (Next n1 k1, Next n2 k2) ->
                if n1 /= n2
                    then False
                    else loop (k1 ()) (k2 ())

main = do
    let t = Node (Node Empty 1 Empty) 2 (Node Empty 3 Empty)
    printNodes t
    print $ addTree t
    
    let tree1 = Node (Node Empty 1 Empty) 2 (Node Empty 3 Empty)
    let tree2 = Node Empty 1 (Node Empty 2 (Node Empty 3 Empty))
    print $ sameFringe tree1 tree2

shiftT :: Monad m => ((a -> ContT r m s) -> ContT s m s) -> ContT s m a
shiftT e = ContT $ \k -> e (lift . k) `runContT` return
 
resetT :: Monad m => ContT a m a -> ContT r m a
resetT e = lift $ e `runContT` return