{-# LANGUAGE GADTs,                       -- match on Refl for type equality
             ExistentialQuantification    -- forall b ans. Yield ...
#-}
{-|
Description : Internal module for type-safe multi-prompt control
Copyright   : (c) 2020, Microsoft Research; Daan Leijen; Ningning Xie
License     : MIT
Maintainer  : xnning@hku.hk; daan@microsoft.com
Stability   : Experimental

Primitive module that implements type safe multi-prompt control.
Used by the "Control.Ev.Eff" module to implement effect handlers.
-}
module Control.Ev.Ctl(
          -- * Markers
            Marker       -- prompt marker
          , markerEq     -- :: Marker a -> Marker b -> Bool

          -- * Control monad
          , Ctl(Pure)    -- multi-prompt control monad
          , runCtl       -- run the control monad       :: Ctl a -> a
          , prompt       -- install a multi-prompt      :: (Marker a -> Ctl a) -> Ctl a
          , yield        -- yield to a specific prompt  :: Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> Ctl b

          -- * Unsafe primitives for "Control.Ev.Eff"
          , unsafeIO            -- lift IO into Ctl        :: IO a -> Ctl a
          , unsafePromptIORef   -- IORef that gets restored per resumption
          ) where

import Prelude hiding (read,flip)
import Control.Monad( ap, liftM )
import Data.Type.Equality( (:~:)( Refl ) )
import Control.Monad.Primitive

-------------------------------------------------------
-- Assume some way to generate a fresh prompt marker
-- associated with specific answer type.
-------------------------------------------------------
import Unsafe.Coerce    ( unsafeCoerce )
import System.IO.Unsafe ( unsafePerformIO )
import Data.IORef

-- | An abstract prompt marker
data Marker a = Marker !Integer

instance Show (Marker a) where
  show :: Marker a -> String
show (Marker Integer
i) = Integer -> String
forall a. Show a => a -> String
show Integer
i

instance Eq (Marker a) where
  Marker a
m1 == :: Marker a -> Marker a -> Bool
== Marker a
m2  = Marker a -> Marker a -> Bool
forall a b. Marker a -> Marker b -> Bool
markerEq Marker a
m1 Marker a
m2

-- | Compare two markers of different types for equality
markerEq :: Marker a -> Marker b -> Bool
markerEq :: Marker a -> Marker b -> Bool
markerEq (Marker Integer
i) (Marker Integer
j)  = (Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
j)

-- if markers match, their types are the same
mmatch :: Marker a -> Marker b -> Maybe ((:~:) a b)
mmatch :: Marker a -> Marker b -> Maybe (a :~: b)
mmatch (Marker Integer
i) (Marker Integer
j) | Integer
i Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
j  = (a :~: b) -> Maybe (a :~: b)
forall a. a -> Maybe a
Just ((Any :~: Any) -> a :~: b
forall a b. a -> b
unsafeCoerce Any :~: Any
forall k (a :: k). a :~: a
Refl)
mmatch Marker a
_ Marker b
_ = Maybe (a :~: b)
forall a. Maybe a
Nothing

-- global unique counter
{-# NOINLINE unique #-}
unique :: IORef Integer
unique :: IORef Integer
unique = IO (IORef Integer) -> IORef Integer
forall a. IO a -> a
unsafePerformIO (Integer -> IO (IORef Integer)
forall a. a -> IO (IORef a)
newIORef Integer
0)

-- evaluate a action with a fresh marker
{-# NOINLINE freshMarker #-}
freshMarker :: (Marker a -> Ctl a) -> Ctl a
freshMarker :: (Marker a -> Ctl a) -> Ctl a
freshMarker Marker a -> Ctl a
f
  = let m :: Integer
m = IO Integer -> Integer
forall a. IO a -> a
unsafePerformIO (IO Integer -> Integer) -> IO Integer -> Integer
forall a b. (a -> b) -> a -> b
$
            do Integer
i <- IORef Integer -> IO Integer
forall a. IORef a -> IO a
readIORef IORef Integer
unique;
               IORef Integer -> Integer -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Integer
unique (Integer
iInteger -> Integer -> Integer
forall a. Num a => a -> a -> a
+Integer
1);
               Integer -> IO Integer
forall (m :: * -> *) a. Monad m => a -> m a
return Integer
i
    in Integer -> Ctl a -> Ctl a
seq Integer
m (Marker a -> Ctl a
f (Integer -> Marker a
forall a. Integer -> Marker a
Marker Integer
m))


{-|  The Multi Prompt control monad,
with existentials `ans` and `b`: where `ans` is the answer type, i.e. the type of the handler/prompt context,
and `b` the result type of the operation.
-}
data Ctl a = Pure { Ctl a -> a
result :: !a }  -- ^ Pure results (only exported for use in the "Control.Ev.Eff" module)
           | forall ans b.
             Yield{ ()
marker :: !(Marker ans),                 -- ^ prompt marker to yield to (in type context `::ans`)
                    ()
op     :: !((b -> Ctl ans) -> Ctl ans),  -- ^ the final action, just needs the resumption (:: b -> Ctl ans) to be evaluated.
                    ()
cont   :: !(b -> Ctl a)                  -- ^ the (partially) build up resumption; `(b -> Ctl a) :~: (b -> Ctl ans)` by the time we reach the prompt
                  }

-- | @yield m op@ yields to a specific marker and calls @op@ in that context
-- with a /resumption/ @k :: b -> Ctl ans@ that resumes at the original call-site
-- with a result of type @b@. If the marker is no longer in the evaluation context,
-- (i.e. it escaped outside its prompt) the `yield` fails with an @"unhandled operation"@ error.
{-# INLINE yield #-}
yield :: Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> Ctl b
yield :: Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> Ctl b
yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op  = Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl b) -> Ctl b
forall a ans b.
Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl a) -> Ctl a
Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op b -> Ctl b
forall a. a -> Ctl a
Pure

{-# INLINE kcompose #-}
kcompose :: (b -> Ctl c) -> (a -> Ctl b) -> a -> Ctl c      -- Kleisli composition
kcompose :: (b -> Ctl c) -> (a -> Ctl b) -> a -> Ctl c
kcompose b -> Ctl c
g a -> Ctl b
f a
x = case (a -> Ctl b
f a
x) of
                   Pure b
x -> b -> Ctl c
g b
x
                   Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op b -> Ctl b
cont -> Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl c) -> Ctl c
forall a ans b.
Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl a) -> Ctl a
Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op (b -> Ctl c
g (b -> Ctl c) -> (b -> Ctl b) -> b -> Ctl c
forall b c a. (b -> Ctl c) -> (a -> Ctl b) -> a -> Ctl c
`kcompose` b -> Ctl b
cont)

{-# INLINE bind #-}
bind :: Ctl a -> (a -> Ctl b) -> Ctl b
bind :: Ctl a -> (a -> Ctl b) -> Ctl b
bind (Pure a
x) a -> Ctl b
f           = a -> Ctl b
f a
x
bind (Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op b -> Ctl a
cont) a -> Ctl b
f  = Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl b) -> Ctl b
forall a ans b.
Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl a) -> Ctl a
Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op (a -> Ctl b
f (a -> Ctl b) -> (b -> Ctl a) -> b -> Ctl b
forall b c a. (b -> Ctl c) -> (a -> Ctl b) -> a -> Ctl c
`kcompose` b -> Ctl a
cont)  -- keep yielding with an extended continuation

instance Functor Ctl where
  fmap :: (a -> b) -> Ctl a -> Ctl b
fmap  = (a -> b) -> Ctl a -> Ctl b
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM
instance Applicative Ctl where
  pure :: a -> Ctl a
pure  = a -> Ctl a
forall (m :: * -> *) a. Monad m => a -> m a
return
  <*> :: Ctl (a -> b) -> Ctl a -> Ctl b
(<*>) = Ctl (a -> b) -> Ctl a -> Ctl b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap
instance Monad Ctl where
  return :: a -> Ctl a
return a
x  = a -> Ctl a
forall a. a -> Ctl a
Pure a
x
  Ctl a
e >>= :: Ctl a -> (a -> Ctl b) -> Ctl b
>>= a -> Ctl b
f   = Ctl a -> (a -> Ctl b) -> Ctl b
forall a b. Ctl a -> (a -> Ctl b) -> Ctl b
bind Ctl a
e a -> Ctl b
f


-- install a prompt with a unique marker (and handle yields to it)
{-# INLINE mprompt #-}
mprompt :: Marker a -> Ctl a -> Ctl a
mprompt :: Marker a -> Ctl a -> Ctl a
mprompt Marker a
m p :: Ctl a
p@(Pure a
_) = Ctl a
p
mprompt Marker a
m (Yield Marker ans
n (b -> Ctl ans) -> Ctl ans
op b -> Ctl a
cont)
  = let cont' :: b -> Ctl a
cont' b
x = Marker a -> Ctl a -> Ctl a
forall a. Marker a -> Ctl a -> Ctl a
mprompt Marker a
m (b -> Ctl a
cont b
x) in  -- extend the continuation with our own prompt
    case Marker a -> Marker ans -> Maybe (a :~: ans)
forall a b. Marker a -> Marker b -> Maybe (a :~: b)
mmatch Marker a
m Marker ans
n of
      Maybe (a :~: ans)
Nothing   -> Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl a) -> Ctl a
forall a ans b.
Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl a) -> Ctl a
Yield Marker ans
n (b -> Ctl ans) -> Ctl ans
op b -> Ctl a
cont'   -- keep yielding (but with the extended continuation)
      Just a :~: ans
Refl -> (b -> Ctl ans) -> Ctl ans
op b -> Ctl a
b -> Ctl ans
cont'           -- found our prompt, invoke `op`.
                   -- Note: `Refl` proves `a ~ ans` (the existential `ans` in Yield)

-- | Install a /prompt/ with a specific prompt `Marker` to which one can `yield`.
-- This connects creation of a marker with instantiating the prompt. The marker passed
-- to the @action@ argument should not escape the @action@ (but this is not statically checked,
-- only at runtime when `yield`ing to it).
{-# INLINE prompt #-}
prompt :: (Marker a -> Ctl a) -> Ctl a
prompt :: (Marker a -> Ctl a) -> Ctl a
prompt Marker a -> Ctl a
action
  = (Marker a -> Ctl a) -> Ctl a
forall a. (Marker a -> Ctl a) -> Ctl a
freshMarker ((Marker a -> Ctl a) -> Ctl a) -> (Marker a -> Ctl a) -> Ctl a
forall a b. (a -> b) -> a -> b
$ \Marker a
m ->   -- create a fresh marker
    Marker a -> Ctl a -> Ctl a
forall a. Marker a -> Ctl a -> Ctl a
mprompt Marker a
m (Marker a -> Ctl a
action Marker a
m)  -- and install a prompt associated with this marker

-- | Run a control monad. This may fail with an @"unhandled operation"@ error if 
-- there is a `yield` to a marker that escaped its prompt scope.
runCtl :: Ctl a -> a
runCtl :: Ctl a -> a
runCtl (Pure a
x)      = a
x
runCtl (Yield Marker ans
_ (b -> Ctl ans) -> Ctl ans
_ b -> Ctl a
_) = String -> a
forall a. HasCallStack => String -> a
error String
"Unhandled operation"  -- only if marker escapes the scope of the prompt


-------------------------------------------------------
-- IORef's
-------------------------------------------------------

-- | Unsafe `IO` in the `Ctl` monad.
{-# INLINE unsafeIO #-}
unsafeIO :: IO a -> Ctl a
unsafeIO :: IO a -> Ctl a
unsafeIO IO a
io = let x :: a
x = IO a -> a
forall (m :: * -> *) a. PrimBase m => m a -> a
unsafeInlinePrim IO a
io in a -> Ctl a -> Ctl a
seq a
x (a -> Ctl a
forall a. a -> Ctl a
Pure a
x)

-- A special prompt that saves and restores state per resumption
mpromptIORef :: IORef a -> Ctl b -> Ctl b
mpromptIORef :: IORef a -> Ctl b -> Ctl b
mpromptIORef IORef a
r Ctl b
action
  = case Ctl b
action of
      p :: Ctl b
p@(Pure b
_) -> Ctl b
p
      Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op b -> Ctl b
cont
        -> do a
val <- IO a -> Ctl a
forall a. IO a -> Ctl a
unsafeIO (IORef a -> IO a
forall a. IORef a -> IO a
readIORef IORef a
r)                 -- save current value on yielding
              let cont' :: b -> Ctl b
cont' b
x = do IO () -> Ctl ()
forall a. IO a -> Ctl a
unsafeIO (IORef a -> a -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef a
r a
val)  -- restore saved value on resume
                               IORef a -> Ctl b -> Ctl b
forall a b. IORef a -> Ctl b -> Ctl b
mpromptIORef IORef a
r (b -> Ctl b
cont b
x)
              Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl b) -> Ctl b
forall a ans b.
Marker ans -> ((b -> Ctl ans) -> Ctl ans) -> (b -> Ctl a) -> Ctl a
Yield Marker ans
m (b -> Ctl ans) -> Ctl ans
op b -> Ctl b
cont'

-- | Create an `IORef` connected to a prompt. The value of
-- the `IORef` is saved and restored through resumptions.
unsafePromptIORef :: a -> (Marker b -> IORef a -> Ctl b) -> Ctl b
unsafePromptIORef :: a -> (Marker b -> IORef a -> Ctl b) -> Ctl b
unsafePromptIORef a
init Marker b -> IORef a -> Ctl b
action
  = (Marker b -> Ctl b) -> Ctl b
forall a. (Marker a -> Ctl a) -> Ctl a
freshMarker ((Marker b -> Ctl b) -> Ctl b) -> (Marker b -> Ctl b) -> Ctl b
forall a b. (a -> b) -> a -> b
$ \Marker b
m ->
    do IORef a
r <- IO (IORef a) -> Ctl (IORef a)
forall a. IO a -> Ctl a
unsafeIO (a -> IO (IORef a)
forall a. a -> IO (IORef a)
newIORef a
init)
       IORef a -> Ctl b -> Ctl b
forall a b. IORef a -> Ctl b -> Ctl b
mpromptIORef IORef a
r (Marker b -> IORef a -> Ctl b
action Marker b
m IORef a
r)