{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Extensible.Effect
-- Copyright   :  (c) Fumiaki Kinoshita 2018
-- License     :  BSD3
--
-- Maintainer  :  Fumiaki Kinoshita <fumiexcel@gmail.com>
--
-- Name-based extensible effects
-----------------------------------------------------------------------------
module Data.Extensible.Effect (
  -- * Base
  Instruction(..)
  , Eff
  , liftEff
  , liftsEff
  , hoistEff
  , castEff
  -- * Step-wise handling
  , Interpreter(..)
  , handleEff
  -- * Peeling
  , peelEff
  , Rebinder
  , rebindEff0
  , peelEff0
  , rebindEff1
  , peelEff1
  , rebindEff2
  , leaveEff
  , retractEff
  -- * Anonymous actions
  , Action(..)
  , Function
  , runAction
  , (@!?)
  , peelAction
  , peelAction0
  -- * transformers-compatible actions and handlers
  -- ** Reader
  , ReaderEff
  , askEff
  , asksEff
  , localEff
  , runReaderEff
  -- ** State
  , State
  , getEff
  , getsEff
  , putEff
  , modifyEff
  , stateEff
  , runStateEff
  , execStateEff
  , evalStateEff
  -- ** Writer
  , WriterEff
  , writerEff
  , tellEff
  , listenEff
  , passEff
  , runWriterEff
  , execWriterEff
  -- ** Maybe
  , MaybeEff
  , nothingEff
  , runMaybeEff
  -- ** Either
  , EitherEff
  , throwEff
  , catchEff
  , runEitherEff
  , mapLeftEff
  -- ** Iter
  , Identity
  , tickEff
  , runIterEff
  -- ** Cont
  , ContT
  , contEff
  , runContEff
  , callCCEff
  ) where

import Control.Applicative
import Data.Bifunctor (first)
import Control.Monad.Skeleton
import Control.Monad.Trans.State.Strict
import Control.Monad.Trans.Cont (ContT(..))
import Data.Extensible.Field
import Data.Extensible.Inclusion
import Data.Extensible.Internal.Rig
import Data.Extensible.Product
import Data.Extensible.Class
import Data.Kind (Type)
import Data.Functor.Identity
import Data.Type.Equality
import Type.Membership

-- | A unit of named effects. This is a variant of @(':|')@ specialised for
-- 'Type -> Type'.
data Instruction (xs :: [Assoc k (Type -> Type)]) a where
  Instruction :: !(Membership xs kv) -> TargetOf kv a -> Instruction xs a

-- | The extensible operational monad
type Eff xs = Skeleton (Instruction xs)

-- | Lift an instruction onto an 'Eff' action.
liftEff :: forall s t xs a. Lookup xs s t => Proxy s -> t a -> Eff xs a
liftEff p x = liftsEff p x id
{-# INLINE liftEff #-}

-- | Lift an instruction onto an 'Eff' action and apply a function to the result.
liftsEff :: forall s t xs a r. Lookup xs s t
  => Proxy s -> t a -> (a -> r) -> Eff xs r
liftsEff _ x k = boned
  $ Instruction (association :: Membership xs (s ':> t)) x :>>= return . k
{-# INLINE liftsEff #-}

-- | Censor a specific type of effects in an action.
hoistEff :: forall s t xs a. Lookup xs s t => Proxy s -> (forall x. t x -> t x) -> Eff xs a -> Eff xs a
hoistEff _ f = hoistSkeleton $ \(Instruction i t) -> case compareMembership (association :: Membership xs (s ':> t)) i of
  Right Refl -> Instruction i (f t)
  _ -> Instruction i t
{-# INLINABLE hoistEff #-}

-- | Upcast an action.
castEff :: IncludeAssoc ys xs => Eff xs a -> Eff ys a
castEff = hoistSkeleton
  $ \(Instruction i t) -> Instruction (hlookup i inclusionAssoc) t

-- | Build a relay-style handler from a triple of functions.
--
-- @
-- runStateEff = peelEff1 (\a s -> return (a, s))
--   (\m k s -> let (a, s') = runState m s in k a s')
-- @
--
peelEff :: forall k t xs a r
  . Rebinder xs r -- ^ Re-bind an unrelated action
  -> (a -> r) -- ^ return the result
  -> (forall x. t x -> (x -> r) -> r) -- ^ Handle the foremost type of an action
  -> Eff (k >: t ': xs) a -> r
peelEff pass ret wrap = go where
  go m = case debone m of
    Return a -> ret a
    Instruction i t :>>= k -> testMembership i
      (\Refl -> wrap t (go . k))
      (\j -> pass (Instruction j t) (go . k))
{-# INLINE peelEff #-}

-- | 'peelEff' specialised for continuations with no argument
peelEff0 :: forall k t xs a r. (a -> Eff xs r) -- ^ return the result
  -> (forall x. t x -> (x -> Eff xs r) -> Eff xs r) -- ^ Handle the foremost type of an action
  -> Eff (k >: t ': xs) a -> Eff xs r
peelEff0 = peelEff rebindEff0
{-# INLINE peelEff0 #-}

-- | 'peelEff' specialised for 1-argument continuation
peelEff1 :: forall k t xs a b r. (a -> b -> Eff xs r) -- ^ return the result
  -> (forall x. t x -> (x -> b -> Eff xs r) -> b -> Eff xs r) -- ^ Handle the foremost type of an action
  -> Eff (k >: t ': xs) a -> b -> Eff xs r
peelEff1 = peelEff rebindEff1
{-# INLINE peelEff1 #-}

-- | A function to bind an 'Instruction' in 'peelEff'.
type Rebinder xs r = forall x. Instruction xs x -> (x -> r) -> r

-- | A common value for the second argument of 'peelEff'. Binds an instruction
-- directly.
rebindEff0 :: Rebinder xs (Eff xs r)
rebindEff0 i k = boned $ i :>>= k

-- | A pre-defined value for the second argument of 'peelEff'.
-- Preserves the argument of the continuation.
rebindEff1 :: Rebinder xs (a -> Eff xs r)
rebindEff1 i k a = boned $ i :>>= flip k a

-- | A pre-defined value for the second argument of 'peelEff'.
-- Preserves two arguments of the continuation.
rebindEff2 :: Rebinder xs (a -> b -> Eff xs r)
rebindEff2 i k a b = boned $ i :>>= \x -> k x a b

-- | Reveal the final result of 'Eff'.
leaveEff :: Eff '[] a -> a
leaveEff m = case debone m of
  Return a -> a
  _ -> error "Impossible"

-- | Tear down an action using the 'Monad' instance of the instruction.
retractEff :: forall k m a. Monad m => Eff '[k >: m] a -> m a
retractEff m = case debone m of
  Return a -> return a
  Instruction i t :>>= k -> testMembership i
    (\Refl -> t >>= retractEff . k)
    $ error "Impossible"

-- | Transformation between effects
newtype Interpreter f g = Interpreter { runInterpreter :: forall a. g a -> f a }

-- | Process an 'Eff' action using a record of 'Interpreter's.
handleEff :: RecordOf (Interpreter m) xs -> Eff xs a -> MonadView m (Eff xs) a
handleEff hs m = case debone m of
  Instruction i t :>>= k -> views (pieceAt i) (\(Field (Interpreter x)) -> x) hs t :>>= k
  Return a -> Return a

-- | Anonymous representation of instructions.
data Action (args :: [Type]) a r where
  AResult :: Action '[] a a
  AArgument :: x -> Action xs a r -> Action (x ': xs) a r

-- | @'Function' [a, b, c] r@ is @a -> b -> c -> r@
type family Function args r :: Type where
  Function '[] r = r
  Function (x ': xs) r = x -> Function xs r

-- | Pass the arguments of 'Action' to the supplied function.
runAction :: Function xs (f a) -> Action xs a r -> f r
runAction r AResult = r
runAction f (AArgument x a) = runAction (f x) a

-- | Create a 'Field' of a 'Interpreter' for an 'Action'.
(@!?) :: FieldName k -> Function xs (f a) -> Field (Interpreter f) (k ':> Action xs a)
_ @!? f = Field $ Interpreter $ runAction f
infix 1 @!?

-- | Specialised version of 'peelEff' for 'Action's.
-- You can pass a function @a -> b -> ... -> (q -> r) -> r@ as a handler for
-- @'Action' '[a, b, ...] q@.
peelAction :: forall k ps q xs a r
  . (forall x. Instruction xs x -> (x -> r) -> r) -- ^ Re-bind an unrelated action
  -> (a -> r) -- ^ return the result
  -> Function ps ((q -> r) -> r) -- ^ Handle the foremost action
  -> Eff (k >: Action ps q ': xs) a -> r
peelAction pass ret wrap = go where
  go m = case debone m of
    Return a -> ret a
    Instruction i t :>>= k -> testMembership i
      (\Refl -> case t of
        (_ :: Action ps q x) ->
          let run :: forall t. Function t ((q -> r) -> r) -> Action t q x -> r
              run f AResult = f (go . k)
              run f (AArgument x a) = run (f x) a
          in run wrap t)
      $ \j -> pass (Instruction j t) (go . k)
{-# INLINE peelAction #-}

-- | Non continuation-passing variant of 'peelAction'.
peelAction0 :: forall k ps q xs a. Function ps (Eff xs q) -- ^ Handle the foremost action
  -> Eff (k >: Action ps q ': xs) a -> Eff xs a
peelAction0 wrap = go where
  go m = case debone m of
    Return a -> return a
    Instruction i t :>>= k -> testMembership i
      (\Refl -> case t of
        (_ :: Action ps q x) ->
          let run :: forall t. Function t (Eff xs q) -> Action t q x -> Eff xs a
              run f AResult = f >>= go . k
              run f (AArgument x a) = run (f x) a
          in run wrap t)
      $ \j -> rebindEff0 (Instruction j t) (go . k)
{-# INLINE peelAction0 #-}

-- | The reader monad is characterised by a type equality between the result
-- type and the enviroment type.
type ReaderEff = (:~:)

-- | Fetch the environment.
askEff :: forall k r xs. Lookup xs k (ReaderEff r)
  => Proxy k -> Eff xs r
askEff p = liftEff p Refl
{-# INLINE askEff #-}

-- | Pass the environment to a function.
asksEff :: forall k r xs a. Lookup xs k (ReaderEff r)
  => Proxy k -> (r -> a) -> Eff xs a
asksEff p = liftsEff p Refl
{-# INLINE asksEff #-}

-- | Modify the enviroment locally.
localEff :: forall k r xs a. Lookup xs k (ReaderEff r)
  => Proxy k -> (r -> r) -> Eff xs a -> Eff xs a
localEff _ f = go where
  go m = case debone m of
    Return a -> return a
    Instruction i t :>>= k -> case compareMembership
      (association :: Membership xs (k >: ReaderEff r)) i of
        Left _ -> boned $ Instruction i t :>>= go . k
        Right Refl -> case t of
          Refl -> boned $ Instruction i t :>>= go . k . f
{-# INLINE localEff #-}

-- | Run the frontal reader effect.
runReaderEff :: forall k r xs a. Eff (k >: ReaderEff r ': xs) a -> r -> Eff xs a
runReaderEff m r = peelEff0 return (\Refl k -> k r) m
{-# INLINE runReaderEff #-}

-- | Get the current state.
getEff :: forall k s xs. Lookup xs k (State s)
  => Proxy k -> Eff xs s
getEff k = liftEff k get
{-# INLINE getEff #-}

-- | Pass the current state to a function.
getsEff :: forall k s a xs. Lookup xs k (State s)
  => Proxy k -> (s -> a) -> Eff xs a
getsEff k = liftsEff k get
{-# INLINE getsEff #-}

-- | Replace the state with a new value.
putEff :: forall k s xs. Lookup xs k (State s)
  => Proxy k -> s -> Eff xs ()
putEff k = liftEff k . put
{-# INLINE putEff #-}

-- | Modify the state.
modifyEff :: forall k s xs. Lookup xs k (State s)
  => Proxy k -> (s -> s) -> Eff xs ()
modifyEff k f = liftEff k $ state $ \s -> ((), f s)
{-# INLINE modifyEff #-}

-- | Lift a state modification function.
stateEff :: forall k s xs a. Lookup xs k (State s)
  => Proxy k -> (s -> (a, s)) -> Eff xs a
stateEff k = liftEff k . state
{-# INLINE stateEff #-}

contState :: State s a -> (a -> s -> r) -> s -> r
contState m k s = let (a, s') = runState m s in k a $! s'

-- | Run the frontal state effect.
runStateEff :: forall k s xs a. Eff (k >: State s ': xs) a -> s -> Eff xs (a, s)
runStateEff = peelEff1 (\a s -> return (a, s)) contState
{-# INLINE runStateEff #-}

-- | Run the frontal state effect and return the final state.
execStateEff :: forall k s xs a. Eff (k >: State s ': xs) a -> s -> Eff xs s
execStateEff = peelEff1 (const return) contState
{-# INLINE execStateEff #-}

-- | Run the frontal state effect and return the final result.
evalStateEff :: forall k s xs a. Eff (k >: State s ': xs) a -> s -> Eff xs a
evalStateEff = peelEff1 (const . return) contState
{-# INLINE evalStateEff #-}

-- | @(,)@ already is a writer monad.
type WriterEff w = (,) w

-- | Write the second element and return the first element.
writerEff :: forall k w xs a. (Lookup xs k (WriterEff w))
  => Proxy k -> (a, w) -> Eff xs a
writerEff k (a, w) = liftEff k (w, a)
{-# INLINE writerEff #-}

-- | Write a value.
tellEff :: forall k w xs. (Lookup xs k (WriterEff w))
  => Proxy k -> w -> Eff xs ()
tellEff k w = liftEff k (w, ())
{-# INLINE tellEff #-}

-- | Squash the outputs into one step and return it.
listenEff :: forall k w xs a. (Lookup xs k (WriterEff w), Monoid w)
  => Proxy k -> Eff xs a -> Eff xs (a, w)
listenEff p = go mempty where
  go w m = case debone m of
    Return a -> writerEff p ((a, w), w)
    Instruction i t :>>= k -> case compareMembership (association :: Membership xs (k ':> (,) w)) i of
      Left _ -> boned $ Instruction i t :>>= go w . k
      Right Refl -> let (w', a) = t
                        !w'' = mappend w w' in go w'' (k a)
{-# INLINE listenEff #-}

-- | Modify the output using the function in the result.
passEff :: forall k w xs a. (Lookup xs k (WriterEff w), Monoid w)
  => Proxy k -> Eff xs (a, w -> w) -> Eff xs a
passEff p = go mempty where
  go w m = case debone m of
    Return (a, f) -> writerEff p (a, f w)
    Instruction i t :>>= k -> case compareMembership (association :: Membership xs (k ':> (,) w)) i of
      Left _ -> boned $ Instruction i t :>>= go w . k
      Right Refl -> let (w', a) = t
                        !w'' = mappend w w' in go w'' (k a)
{-# INLINE passEff #-}

contWriter :: Monoid w => (w, a) -> (a -> w -> r) -> w -> r
contWriter (w', a) k w = k a $! mappend w w'

-- | Run the frontal writer effect.
runWriterEff :: forall k w xs a. Monoid w => Eff (k >: WriterEff w ': xs) a -> Eff xs (a, w)
runWriterEff = peelEff1 (\a w -> return (a, w)) contWriter `flip` mempty
{-# INLINE runWriterEff #-}

-- | Run the frontal state effect.
execWriterEff :: forall k w xs a. Monoid w => Eff (k >: WriterEff w ': xs) a -> Eff xs w
execWriterEff = peelEff1 (const return) contWriter `flip` mempty
{-# INLINE execWriterEff #-}

-- | An effect with no result
type MaybeEff = Const ()

-- | Break out of the computation. Similar to 'Nothing'.
nothingEff :: Lookup xs k MaybeEff => Proxy k -> Eff xs a
nothingEff = flip throwEff ()

-- | Run an effect which may fail in the name of @k@.
runMaybeEff :: forall k xs a. Eff (k >: MaybeEff ': xs) a -> Eff xs (Maybe a)
runMaybeEff = peelEff0 (return . Just) $ \_ _ -> return Nothing
{-# INLINE runMaybeEff #-}

-- | Throwing an exception
type EitherEff = Const

-- | Throw an exception @e@, throwing the rest of the computation away.
throwEff :: Lookup xs k (EitherEff e) => Proxy k -> e -> Eff xs a
throwEff k = liftEff k . Const
{-# INLINE throwEff #-}

-- | Attach a handler for an exception.
catchEff :: forall k e xs a. (Lookup xs k (EitherEff e))
  => Proxy k -> Eff xs a -> (e -> Eff xs a) -> Eff xs a
catchEff _ m0 handler = go m0 where
  go m = case debone m of
    Return a -> return a
    Instruction i t :>>= k -> case compareMembership (association :: Membership xs (k ':> Const e)) i of
      Left _ -> boned $ Instruction i t :>>= go . k
      Right Refl -> handler (getConst t)
{-# INLINE catchEff #-}

-- | Run an action and abort on 'throwEff'.
runEitherEff :: forall k e xs a. Eff (k >: EitherEff e ': xs) a -> Eff xs (Either e a)
runEitherEff = peelEff0 (return . Right) $ \(Const e) _ -> return $ Left e
{-# INLINE runEitherEff #-}

-- | Put a milestone on a computation.
tickEff :: Lookup xs k Identity => Proxy k -> Eff xs ()
tickEff k = liftEff k $ Identity ()
{-# INLINE tickEff #-}

mapHeadEff :: (forall x. s x -> t x) -> Eff ((k >: s) ': xs) a -> Eff ((k' >: t) ': xs) a
mapHeadEff f = hoistSkeleton $ \(Instruction i t) -> testMembership i
  (\Refl -> Instruction leadership $ f t)
  (\j -> Instruction (nextMembership j) t)

-- | Take a function and applies it to an Either effect iff the effect takes the form Left _.
mapLeftEff :: (e -> e') -> Eff ((k >: EitherEff e) ': xs) a -> Eff ((k >: EitherEff e') ': xs) a
mapLeftEff f = mapHeadEff (first f)

-- | Run a computation until the first call of 'tickEff'.
runIterEff :: Eff (k >: Identity ': xs) a
  -> Eff xs (Either a (Eff (k >: Identity ': xs) a))
runIterEff m = case debone m of
  Return a -> return (Left a)
  Instruction i t :>>= k -> testMembership i
    (\Refl -> return $ Right $ k $ runIdentity t)
    $ \j -> boned $ Instruction j t :>>= runIterEff . k

-- | Place a continuation-passing action.
contEff :: Lookup xs k (ContT r m) => Proxy k
  -> ((a -> m r) -> m r) -> Eff xs a
contEff k = liftEff k . ContT

-- | Unwrap a continuation.
runContEff :: forall k r xs a. Eff (k >: ContT r (Eff xs) ': xs) a
  -> (a -> Eff xs r)
  -> Eff xs r
runContEff m cont = case debone m of
  Return a -> cont a
  Instruction i t :>>= k -> testMembership i
    (\Refl -> runContT t (flip runContEff cont . k))
    $ \j -> boned $ Instruction j t :>>= flip runContEff cont . k

-- | Call a function with the current continuation as its argument
callCCEff :: Proxy k -> ((a -> Eff ((k >: ContT r (Eff xs)) : xs) b) -> Eff ((k >: ContT r (Eff xs)) : xs) a) -> Eff ((k >: ContT r (Eff xs)) : xs) a
callCCEff k f = contHead k . ContT $ \c -> runContEff (f (\x -> contHead k . ContT $ \_ -> c x)) c
  where
    contHead :: Proxy k -> ContT r (Eff xs) a -> Eff ((k >: ContT r (Eff xs)) ': xs) a
    contHead _ c = boned $ Instruction leadership c :>>= return