Continuations

Continuations.hs

{-# LANGUAGE Rank2Types, MultiParamTypeClasses, FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts, TypeApplications, RecordWildCards #-}
{-# LANGUAGE UndecidableInstances, TypeFamilies, ExistentialQuantification #-}
{-# LANGUAGE ScopedTypeVariables, LambdaCase #-}
{-# OPTIONS_GHC -Wall #-}
module Continuations where
import Control.Monad.Trans
import Control.Monad.Base
import Control.Monad.Trans.Control
import Control.Monad.State.Lazy
import Control.Monad.Reader

import Unsafe.Coerce

newtype ContT f m a = ContT { runContT :: forall b. (a -> m (f b)) -> m (f b) }

instance (MonadBase b m) => MonadBase b (ContT f m) where
    liftBase m = lift (liftBase m)

handleContT :: (Applicative f, Monad m) => ContT f m a -> m (f a)
handleContT c = runContT c (return . pure)

instance Functor (ContT f m) where
    fmap f (ContT g) = ContT $ \amb -> g (amb . f)

instance MonadTrans (ContT f) where
    lift m = ContT $ \f -> m >>= f

instance Monad (ContT f m) where
    return a = ContT $ \f -> f a
    ContT f >>= g = ContT $ \h -> f (flip runContT h . g)

instance Applicative (ContT f m) where
    pure = return
    f <*> a = f >>= (`fmap` a)

instance (MonadIO m) => MonadIO (ContT f m) where
    liftIO i = lift (liftIO i)

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

instance (MonadReader r m, Traversable f, Monad f) => MonadReader r (ContT f m) where
    ask = lift ask
    local f = (>>= choose) . lift . local f . handleContT

callCC :: (Monad m, Monad f) => (forall b. (a -> ContT f m (f b)) -> ContT f m (f b)) -> ContT f m a
callCC f = ContT $ \k -> join <$> handleContT (f $ lift . k)

choose :: (Monad m, Traversable f, Monad f) => f a -> ContT f m a
choose l = callCC $ fmap join . forM l

liftWithF :: (Traversable f, Monad f, Monad m) => m (f a) -> ContT f m a
liftWithF m = lift m >>= choose

instance (Traversable f, Monad f) => MonadTransControl (ContT f) where
    type StT (ContT f) a = f a
    liftWith f = lift (f handleContT)
    restoreT = (>>= choose) . lift

data SelfCont m = forall a. SelfCont (ReaderT SearchScheme (StateT [SelfCont m] m) (Maybe a))

class Search m where
    search :: [a] -> m a

data SearchScheme = BFS | DFS

instance
    Monad m
    => Search (ContT Maybe (ReaderT SearchScheme (StateT [SelfCont m] m))) where
    search as = ContT $ \k -> do
        ss <- ask
        let add = case ss of BFS -> flip (++); DFS -> (++)
        modify' (add (map (SelfCont . k) as))
        return Nothing

maybeToList = maybe [] return

exhaust :: forall a m. Monad m => ReaderT SearchScheme (StateT [SelfCont m] m) [a]
exhaust = do
    conts <- get
    case conts of
        [] -> return []
        (SelfCont x : xs) -> do
            put xs
            as <- unsafeCoerce @_ @(Maybe a) <$> x
            rest <- exhaust
            return (maybeToList as ++ rest)

runSearch :: forall m a. Monad m
        => SearchScheme -> ContT Maybe (ReaderT SearchScheme (StateT [SelfCont m] m)) a -> m [a]
runSearch ss m = evalStateT (runReaderT (handleContT m >> exhaust) ss) []

test :: (Search m, MonadIO m) => m Int
test = do
    n <- search [1, 2]
    liftIO (putStrLn $ "n: " ++ show n)
    m <- search [1, 2]
    liftIO (putStrLn $ "m: " ++ show m)
    o <- search [1, 2]
    liftIO (putStrLn $ "o: " ++ show o)
    return (n + m + o)