{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE UndecidableInstances #-}

module DelCont (DelCont, DelContT, MonadDelCont(..), runDelCont, runDelContT) where

import Control.Monad
import Control.Monad.Trans
import Control.Monad.Identity

import Control.Monad.State
import Control.Monad.Reader
import Control.Monad.Writer

import Seq
import qualified Prompt

-- The frame-based implementation from the Dybvig, Peyton-Jones, Sabry paper.
-- It was the one of the three that was called 'properly tail-recursive,' which
-- I expect is a good property.
newtype Frame m r a b = Frame (a -> DelContT r m b)
type Cont r m a b = Seq (Frame m) r a b
newtype SubCont r m a b = SC (Seq (Frame m) r a b)

-- The two exported types. Delimited continuation monad and transformer
-- I changed the names from CC to DelCont, as that seemed like a more MTL-ish
-- name.
newtype DelContT r m a = DCT (forall c. Cont r m a c -> Prompt.P r m c)

newtype DelCont r a = DelCont { unDelCont :: DelContT r Identity a }
    deriving (Functor, Monad, MonadDelCont (Prompt.Prompt r) (SubCont r Identity))

runDelCont :: (forall r. DelCont r a) -> a
runDelCont dc = runIdentity (runDelContT (unDelCont dc))

appk :: (Monad m) => Cont r m a c -> a -> Prompt.P r m c
appk EmptyS v = return v
appk (PushP _ k) v = appk k v
appk (PushSeg (Frame f) k) v = let DCT e = f v in e k
appk (PushCO f k) v = appk k (f v)

runDelContT :: (Monad m) => (forall r. DelContT r m a) -> m a
runDelContT dc = Prompt.runP (let DCT e = dc in e EmptyS)

-- The new typeclass for monad stacks that allow delimited control
class (Monad m) => MonadDelCont p s m | m -> p s where
    newPrompt   :: m (p a)
    pushPrompt  :: p a -> m a -> m a
    withSubCont :: p b -> (s a b -> m b) -> m a
    pushSubCont :: s a b -> m a -> m b

-- DelContT instances for various Monad classes, including, of course,
-- MonadDelCont
instance (Monad m) => MonadDelCont (Prompt.Prompt r) (SubCont r m) (DelContT r m) where
    newPrompt = DCT (\k -> Prompt.newPrompt >>= appk k)
    pushPrompt p (DCT e) = DCT (\k -> e (PushP p k))
    withSubCont p f = DCT (\k -> let (subk, k') = splitSeq p k
                                     DCT e = f (SC subk)
                                  in e k')
    pushSubCont (SC subk) (DCT e) = DCT (\k -> e (appendSeq subk k))

-- This instance was cribbed from the implementation of the zipper OS here:
-- http://okmij.org/ftp/Computation/Continuations.html#zipper-fs
instance MonadTrans (DelContT r) where
    lift m = DCT (\k -> lift m >>= appk k)

instance (Monad m) => Monad (DelContT r m) where
    return v = DCT (\k -> appk k v)
    (DCT e) >>= f = DCT (\k -> e (PushSeg (Frame f) k))

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

instance (MonadReader e m) => MonadReader e (DelContT r m) where
    ask = lift ask
    local f (DCT dc) = DCT (\k -> local f (dc k))

-- Instances of MonadDelCont for various other monad transformers.
-- Requires transforming prompts and subcontinuations to fit the 'hidden'
-- modifications of the monad transformers.

newtype PairedPrompt s p a = PP { unPP :: p (a, s) }
newtype PairedSubCont s k a b = PSC { unPSC :: k (a, s) (b, s) }

instance (MonadDelCont p s m) =>
            MonadDelCont (PairedPrompt t p) (PairedSubCont t s) (StateT t m) where
    newPrompt = PP `liftM` lift newPrompt
    pushPrompt p m = StateT $ \s -> pushPrompt (unPP p) (runStateT m s)
    withSubCont p f = StateT $ \s ->
                        withSubCont (unPP p) (\k -> runStateT (f (PSC k)) s)
    pushSubCont k m = StateT $ \s ->
                        pushSubCont (unPSC k) (runStateT m s)

instance (MonadDelCont p s m) => MonadDelCont p s (ReaderT r m) where
    newPrompt = lift newPrompt
    pushPrompt p m = ReaderT $ \r -> pushPrompt p (runReaderT m r)
    withSubCont p f = ReaderT $ \r -> withSubCont p (\k -> runReaderT (f k) r)
    pushSubCont k m = ReaderT $ \r -> pushSubCont k (runReaderT m r)

instance (MonadDelCont p s m, Monoid w) => 
            MonadDelCont (PairedPrompt w p) (PairedSubCont w s) (WriterT w m) where
    newPrompt = PP `liftM` lift newPrompt
    pushPrompt p m = WriterT $ pushPrompt (unPP p) (runWriterT m)
    withSubCont p f = WriterT $
                        withSubCont (unPP p) (\k -> runWriterT (f (PSC k)))
    pushSubCont k m = WriterT $ pushSubCont (unPSC k) (runWriterT m)