Cont.hs

{-# LANGUAGE DeriveFunctor, FlexibleContexts, FlexibleInstances, KindSignatures, ExistentialQuantification, RankNTypes, StandaloneDeriving, GeneralizedNewtypeDeriving, TypeOperators, MultiParamTypeClasses, QuantifiedConstraints, UndecidableInstances #-}

module Control.Effect.Cont where

import Control.Effect.Carrier
import Control.Monad
import Control.Effect.Sum
import Data.Coerce

data Cont (ref :: * -> *) (m :: * -> *) k
  = forall a . Jump  (ref a) a
  | forall a . Subst (ref a -> k) (a -> k)

instance HFunctor (Cont ref) where
  hmap _ = coerce

deriving instance Functor (Cont ref m)

jump :: (Member (Cont (ref m)) sig, Carrier sig m) => ref m a -> a -> m b
jump r v = send (Jump r v)

subst :: (Member (Cont (Ref m)) sig, Carrier sig m) => (Ref m a -> m b) -> (a -> m b) -> m b
subst r alt = send (Subst r alt)

callCC :: (Member (Cont (Ref m)) sig, Carrier sig m) => ((b -> m b) -> m b) -> m b
callCC f = subst (f . jump) pure

-- | The continuation monad, as per 'ContT'. A 'MonadTrans' instance is omitted
-- because this is not a functor in the category of monads. This should probably
-- always be the topmost carrier in an effect stack, lest you anger the gods.
newtype ContC r m a = ContC { runContC :: (a -> m r) -> m r }

instance Functor (ContC r m) where
  fmap f m = ContC $ \ c -> runContC m (c . f)
  {-# INLINE fmap #-}

instance Applicative (ContC r m) where
  pure x  = ContC ($ x)
  {-# INLINE pure #-}
  f <*> v = ContC $ \ c -> runContC f $ \ g -> runContC v (c . g)
  {-# INLINE (<*>) #-}
  m *> k = m >>= \_ -> k
  {-# INLINE (*>) #-}

instance Monad (ContC r m) where
  m >>= k  = ContC $ \ c -> runContC m (\ x -> runContC (k x) c)
  {-# INLINE (>>=) #-}

evalContC :: Applicative m => ContC r m r -> m r
evalContC m = runContC m pure

-- Copied from Control.Monad.Cont
primCallCC :: ((a -> ContC r m b) -> ContC r m a) -> ContC r m a
primCallCC f = ContC $ \ c -> runContC (f (\ x -> ContC $ \ _ -> c x)) c

data Ref m a = forall r . Ref { unRef :: a -> m r }

instance (Effect sig, Carrier sig m) => Carrier (Cont (Ref (ContC r m)) :+: sig) (ContC r m) where
  eff (R other) = ContC $ \x -> eff $ hmap (error "[BUG]: ContC reached an unreachable (?) hmap") $ fmap' (flip runContC x) other
  eff (L x) = ContC . runContC $ case x of
    Jump (Ref exit) v -> exit v *> pure (error "[BUG]: ContC Jump didn't terminate properly")
    Subst main eh -> primCallCC $ \exit -> main (Ref (eh >=> exit))