{-|
Description : Definitions for some common effects.
Copyright   : (c) 2021, Microsoft Research; Daan Leijen; Ningning Xie
License     : MIT
Maintainer  : xnning@hku.hk; daan@microsoft.com
Stability   : Experimental

Some definitions for common effects.
-}
{-# LANGUAGE TypeOperators, FlexibleContexts, FlexibleInstances, MultiParamTypeClasses, UndecidableInstances, Rank2Types #-}
module Control.Mp.Util
  (
    -- * Reader
    Reader(Reader,ask)
  , reader
    -- * State
  , State(State,get,put)
  , state
    -- * Writer
  , Writer(Writer,tell)
  , writer
    -- * Exception
  , Except(Except,throwError)
  , catchError, exceptEither, exceptMaybe, exceptDefault
    -- * Choice
  , Choose(Choose,none,choose)
  , chooseFirst, chooseAll
) where

import Control.Mp.Eff
import Control.Monad
import Control.Applicative

------------
-- Reader
------------

-- | A standard reader effect for values of type @a@.
data Reader a e ans = Reader { Reader a e ans -> Op () a e ans
ask :: !(Op () a e ans)  -- ^ get the reader value of type @a@ (as @`perform` ask ()@)
                             }

-- | A handler for a `Reader` effect with a value of type @a@.
{-# INLINE reader #-}
reader :: a -> Eff (Reader a :* e) ans -> Eff e ans
reader :: a -> Eff (Reader a :* e) ans -> Eff e ans
reader a
x
  = Reader a e ans -> Eff (Reader a :* e) ans -> Eff e ans
forall (h :: * -> * -> *) e ans.
h e ans -> Eff (h :* e) ans -> Eff e ans
handler (Reader :: forall a e ans. Op () a e ans -> Reader a e ans
Reader{ ask :: Op () a e ans
ask = a -> Op () a e ans
forall a e ans. a -> Op () a e ans
value a
x })

{-
-- does not work due to the functional dependency in MonadReader
instance (Reader a :? e) => MR.MonadReader a (Eff e) where
  ask       = perform ask ()
-}

------------
-- State
------------

-- | A standard state effect of type @a@.
data State a e ans = State { State a e ans -> Op () a e ans
get :: !(Op () a e ans) -- ^ Get the current state (as @`perform` get ()@)
                           , State a e ans -> Op a () e ans
put :: !(Op a () e ans) -- ^ Set the current state (as @`perform` put x@)
                           }

-- | A state handler that takes an initial state of type @a@.
{-# INLINE state #-}
state :: a -> Eff (State a :* e) ans -> Eff e ans
state :: a -> Eff (State a :* e) ans -> Eff e ans
state a
init
  = a
-> State a (Local a :* e) ans
-> Eff (State a :* e) ans
-> Eff e ans
forall a (h :: * -> * -> *) e ans.
a -> h (Local a :* e) ans -> Eff (h :* e) ans -> Eff e ans
handlerLocal a
init (State :: forall a e ans. Op () a e ans -> Op a () e ans -> State a e ans
State{ get :: Op () a (Local a :* e) ans
get = (() -> Eff (Local a :* e) a) -> Op () a (Local a :* e) ans
forall a e b ans. (a -> Eff e b) -> Op a b e ans
function (\()
_ -> Eff (Local a :* e) a
forall a e. Eff (Local a :* e) a
localGet),
                              put :: Op a () (Local a :* e) ans
put = (a -> Eff (Local a :* e) ()) -> Op a () (Local a :* e) ans
forall a e b ans. (a -> Eff e b) -> Op a b e ans
function (\a
x -> a -> Eff (Local a :* e) ()
forall a e. a -> Eff (Local a :* e) ()
localPut a
x) })


{-
-- does not work due to the functional dependency in MonadState
instance (State a :? e) => MS.MonadState a (Eff e) where
  get   = perform get ()
  put x = perform put x
-}

------------
-- Writer
------------

-- | A standard writer effect of type @a@
data Writer a e ans = Writer { Writer a e ans -> Op a () e ans
tell :: !(Op a () e ans) -- ^ Output a value of type @a@ (as @`perform` tell msg@)
                             }

-- | A standard `Writer` handler for any monoidal type @a@. Returns the final
-- result of type @ans@ and the appended @tell@ arguments.
{-# INLINE writer #-}
writer :: (Monoid a) => Eff (Writer a :* e) ans -> Eff e (ans,a)
writer :: Eff (Writer a :* e) ans -> Eff e (ans, a)
writer
  = a
-> (ans -> a -> (ans, a))
-> Writer a (Local a :* e) (ans, a)
-> Eff (Writer a :* e) ans
-> Eff e (ans, a)
forall a ans b (h :: * -> * -> *) e.
a
-> (ans -> a -> b)
-> h (Local a :* e) b
-> Eff (h :* e) ans
-> Eff e b
handlerLocalRet a
forall a. Monoid a => a
mempty (,) (Writer a (Local a :* e) (ans, a)
 -> Eff (Writer a :* e) ans -> Eff e (ans, a))
-> Writer a (Local a :* e) (ans, a)
-> Eff (Writer a :* e) ans
-> Eff e (ans, a)
forall a b. (a -> b) -> a -> b
$
    Writer :: forall a e ans. Op a () e ans -> Writer a e ans
Writer{ tell :: Op a () (Local a :* e) (ans, a)
tell = (a -> Eff (Local a :* e) ()) -> Op a () (Local a :* e) (ans, a)
forall a e b ans. (a -> Eff e b) -> Op a b e ans
function (\a
x -> do{ a
xs <- Eff (Local a :* e) a
forall a e. Eff (Local a :* e) a
localGet; a -> Eff (Local a :* e) ()
forall a e. a -> Eff (Local a :* e) ()
localPut (a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
xs a
x); () -> Eff (Local a :* e) ()
forall (m :: * -> *) a. Monad m => a -> m a
return () }) }


------------
-- Except
------------

-- | A standard exception effect, throwing values of type @a@.
data Except a e ans = Except { Except a e ans -> forall b. Op a b e ans
throwError :: !(forall b. Op a b e ans) -- ^ Throw an exception with a value of type @a@ (as @`perform` throwError x@)
                             }

-- | Handle an exception.
catchError :: Eff (Except a :* e) ans -> (a -> Eff e ans) -> Eff e ans
catchError :: Eff (Except a :* e) ans -> (a -> Eff e ans) -> Eff e ans
catchError Eff (Except a :* e) ans
action a -> Eff e ans
h
  = Except a e ans -> Eff (Except a :* e) ans -> Eff e ans
forall (h :: * -> * -> *) e ans.
h e ans -> Eff (h :* e) ans -> Eff e ans
handler (Except :: forall a e ans. (forall b. Op a b e ans) -> Except a e ans
Except{ throwError :: forall b. Op a b e ans
throwError = (a -> Eff e ans) -> Op a b e ans
forall a e ans b. (a -> Eff e ans) -> Op a b e ans
except (\a
x -> a -> Eff e ans
h a
x) }) Eff (Except a :* e) ans
action

-- | Transform an exception effect to an @Either@ type.
exceptEither :: Eff (Except a :* e) ans -> Eff e (Either a ans)
exceptEither :: Eff (Except a :* e) ans -> Eff e (Either a ans)
exceptEither
  = (ans -> Either a ans)
-> Except a e (Either a ans)
-> Eff (Except a :* e) ans
-> Eff e (Either a ans)
forall ans a (h :: * -> * -> *) e.
(ans -> a) -> h e a -> Eff (h :* e) ans -> Eff e a
handlerRet ans -> Either a ans
forall a b. b -> Either a b
Right (Except :: forall a e ans. (forall b. Op a b e ans) -> Except a e ans
Except{ throwError :: forall b. Op a b e (Either a ans)
throwError = (a -> Eff e (Either a ans)) -> Op a b e (Either a ans)
forall a e ans b. (a -> Eff e ans) -> Op a b e ans
except (\a
x -> Either a ans -> Eff e (Either a ans)
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Either a ans
forall a b. a -> Either a b
Left a
x) ) })

-- | Remove the exception effect using a default value in case an exception was thrown.
exceptDefault :: ans -> Eff (Except a :* e) ans -> Eff e ans
exceptDefault :: ans -> Eff (Except a :* e) ans -> Eff e ans
exceptDefault ans
def
  = Except a e ans -> Eff (Except a :* e) ans -> Eff e ans
forall (h :: * -> * -> *) e ans.
h e ans -> Eff (h :* e) ans -> Eff e ans
handler (Except :: forall a e ans. (forall b. Op a b e ans) -> Except a e ans
Except{ throwError :: forall b. Op a b e ans
throwError = (a -> Eff e ans) -> Op a b e ans
forall a e ans b. (a -> Eff e ans) -> Op a b e ans
except (\a
_ -> ans -> Eff e ans
forall (m :: * -> *) a. Monad m => a -> m a
return ans
def) })

-- | Transform an exception effect to a @Maybe@ type.
exceptMaybe :: Eff (Except a :* e) ans -> Eff e (Maybe ans)
exceptMaybe :: Eff (Except a :* e) ans -> Eff e (Maybe ans)
exceptMaybe
  = (ans -> Maybe ans)
-> Except a e (Maybe ans)
-> Eff (Except a :* e) ans
-> Eff e (Maybe ans)
forall ans a (h :: * -> * -> *) e.
(ans -> a) -> h e a -> Eff (h :* e) ans -> Eff e a
handlerRet ans -> Maybe ans
forall a. a -> Maybe a
Just (Except :: forall a e ans. (forall b. Op a b e ans) -> Except a e ans
Except{ throwError :: forall b. Op a b e (Maybe ans)
throwError = (a -> Eff e (Maybe ans)) -> Op a b e (Maybe ans)
forall a e ans b. (a -> Eff e ans) -> Op a b e ans
except (\a
_ -> Maybe ans -> Eff e (Maybe ans)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ans
forall a. Maybe a
Nothing) })



--------------------------------------------------------------------------------
-- Choose
--------------------------------------------------------------------------------

-- | Choose implements backtracking.
data Choose e ans = Choose { Choose e ans -> forall a. Op () a e ans
none   :: !(forall a. Op () a e ans)  -- ^ @`perform none ()` indicates no result
                           , Choose e ans -> Op Int Int e ans
choose :: !(Op Int Int e ans)         -- ^ @`perform` choose n` resumes up to @n@ times (returning @1@ up to @n@)
                           }


-- | Return the first result found in a computation using `choose` for backtracking.
chooseFirst :: Eff (Choose :* e) ans -> Eff e (Maybe ans)
chooseFirst :: Eff (Choose :* e) ans -> Eff e (Maybe ans)
chooseFirst
  = (ans -> Maybe ans)
-> Choose e (Maybe ans)
-> Eff (Choose :* e) ans
-> Eff e (Maybe ans)
forall ans a (h :: * -> * -> *) e.
(ans -> a) -> h e a -> Eff (h :* e) ans -> Eff e a
handlerRet ans -> Maybe ans
forall a. a -> Maybe a
Just (Choose e (Maybe ans)
 -> Eff (Choose :* e) ans -> Eff e (Maybe ans))
-> Choose e (Maybe ans)
-> Eff (Choose :* e) ans
-> Eff e (Maybe ans)
forall a b. (a -> b) -> a -> b
$
    Choose :: forall e ans.
(forall a. Op () a e ans) -> Op Int Int e ans -> Choose e ans
Choose{ none :: forall a. Op () a e (Maybe ans)
none   = (() -> Eff e (Maybe ans)) -> Op () a e (Maybe ans)
forall a e ans b. (a -> Eff e ans) -> Op a b e ans
except (\()
_ -> Maybe ans -> Eff e (Maybe ans)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ans
forall a. Maybe a
Nothing)
          , choose :: Op Int Int e (Maybe ans)
choose = (Int -> (Int -> Eff e (Maybe ans)) -> Eff e (Maybe ans))
-> Op Int Int e (Maybe ans)
forall a b e ans.
(a -> (b -> Eff e ans) -> Eff e ans) -> Op a b e ans
operation (\Int
hi Int -> Eff e (Maybe ans)
k -> let try :: Int -> Eff e (Maybe ans)
try Int
n = if (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
hi)
                                                      then Maybe ans -> Eff e (Maybe ans)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ans
forall a. Maybe a
Nothing
                                                      else do Maybe ans
x <- Int -> Eff e (Maybe ans)
k Int
n
                                                              case Maybe ans
x of
                                                                Maybe ans
Nothing -> Int -> Eff e (Maybe ans)
try (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                                                                Maybe ans
_       -> Maybe ans -> Eff e (Maybe ans)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ans
x
                                         in Int -> Eff e (Maybe ans)
try Int
1)
          }

-- | Return all possible results found in a computation using `choose` for backtracking
chooseAll :: Eff (Choose :* e) a -> Eff e [a]
chooseAll :: Eff (Choose :* e) a -> Eff e [a]
chooseAll
  = (a -> [a]) -> Choose e [a] -> Eff (Choose :* e) a -> Eff e [a]
forall ans a (h :: * -> * -> *) e.
(ans -> a) -> h e a -> Eff (h :* e) ans -> Eff e a
handlerRet (\a
x -> [a
x]) (Choose e [a] -> Eff (Choose :* e) a -> Eff e [a])
-> Choose e [a] -> Eff (Choose :* e) a -> Eff e [a]
forall a b. (a -> b) -> a -> b
$
    Choose :: forall e ans.
(forall a. Op () a e ans) -> Op Int Int e ans -> Choose e ans
Choose{ none :: forall a. Op () a e [a]
none   = (() -> Eff e [a]) -> Op () a e [a]
forall a e ans b. (a -> Eff e ans) -> Op a b e ans
except (\()
_ -> [a] -> Eff e [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [])
          , choose :: Op Int Int e [a]
choose = (Int -> (Int -> Eff e [a]) -> Eff e [a]) -> Op Int Int e [a]
forall a b e ans.
(a -> (b -> Eff e ans) -> Eff e ans) -> Op a b e ans
operation (\Int
hi Int -> Eff e [a]
k -> let collect :: Int -> [a] -> Eff e [a]
collect Int
0 [a]
acc = [a] -> Eff e [a]
forall (m :: * -> *) a. Monad m => a -> m a
return [a]
acc
                                             collect Int
n [a]
acc = do [a]
xs <- Int -> Eff e [a]
k Int
n
                                                                Int -> [a] -> Eff e [a]
collect (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) ([a] -> Eff e [a]) -> [a] -> Eff e [a]
forall a b. (a -> b) -> a -> b
$! ([a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ [a]
acc)
                                         in Int -> [a] -> Eff e [a]
collect Int
hi [])
          }

instance (Choose :? e) => Alternative (Eff e) where
  empty :: Eff e a
empty      = (forall e' ans. Choose e' ans -> Op () a e' ans) -> () -> Eff e a
forall (h :: * -> * -> *) e a b.
In h e =>
(forall e' ans. h e' ans -> Op a b e' ans) -> a -> Eff e b
perform forall e' ans. Choose e' ans -> Op () a e' ans
forall e ans. Choose e ans -> forall a. Op () a e ans
none ()
  Eff e a
m1 <|> :: Eff e a -> Eff e a -> Eff e a
<|> Eff e a
m2  = do Int
x <- (forall e' ans. Choose e' ans -> Op Int Int e' ans)
-> Int -> Eff e Int
forall (h :: * -> * -> *) e a b.
In h e =>
(forall e' ans. h e' ans -> Op a b e' ans) -> a -> Eff e b
perform forall e' ans. Choose e' ans -> Op Int Int e' ans
choose Int
2
                  if (Int
xInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
1) then Eff e a
m1 else Eff e a
m2

instance (Choose :? e) => MonadPlus (Eff e) where
  mzero :: Eff e a
mzero       = (forall e' ans. Choose e' ans -> Op () a e' ans) -> () -> Eff e a
forall (h :: * -> * -> *) e a b.
In h e =>
(forall e' ans. h e' ans -> Op a b e' ans) -> a -> Eff e b
perform forall e' ans. Choose e' ans -> Op () a e' ans
forall e ans. Choose e ans -> forall a. Op () a e ans
none ()
  mplus :: Eff e a -> Eff e a -> Eff e a
mplus Eff e a
m1 Eff e a
m2 = do Int
x <- (forall e' ans. Choose e' ans -> Op Int Int e' ans)
-> Int -> Eff e Int
forall (h :: * -> * -> *) e a b.
In h e =>
(forall e' ans. h e' ans -> Op a b e' ans) -> a -> Eff e b
perform forall e' ans. Choose e' ans -> Op Int Int e' ans
choose Int
2
                   if (Int
xInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
1) then Eff e a
m1 else Eff e a
m2