{-# LANGUAGE Trustworthy #-}
-- |
-- Copyright: (c) 2021 Xy Ren
-- License: BSD3
-- Maintainer: xy.r@outlook.com
-- Stability: experimental
-- Portability: non-portable (GHC only)
module Cleff.State
  ( -- * Effect
    State (..)
    -- * Operations
  , get
  , put
  , state
  , gets
  , modify
    -- * Interpretations
  , runState
  , runStateIORef
  , runStateMVar
  , runStateTVar
  , zoom
  ) where

import           Cleff
import           Cleff.Internal.Base
import           Control.Monad       (void)
import           Data.Atomics        (atomicModifyIORefCAS)
import           Data.Tuple          (swap)
import           Lens.Micro          (Lens', (&), (.~), (^.))
import           UnliftIO.IORef      (IORef, newIORef, readIORef, writeIORef)
import           UnliftIO.MVar       (MVar, modifyMVar, readMVar, swapMVar)
import           UnliftIO.STM        (TVar, atomically, readTVar, readTVarIO, writeTVar)

-- * Effect

-- | An effect capable of providing a mutable state @s@ that can be read and written. This roughly corresponds to the
-- @MonadState@ typeclass and @StateT@ monad transformer in the @mtl@ library.
data State s :: Effect where
  Get :: State s m s
  Put :: s -> State s m ()
  State :: (s -> (a, s)) -> State s m a

-- * Operations

makeEffect_ ''State

-- | Read the current state.
get :: State s :> es => Eff es s

-- | Update the state with a new value.
put :: State s :> es => s -> Eff es ()

-- | Modify the state /and/ produce a value from the state via a function.
state :: State s :> es
  => (s -> (a, s)) -- ^ The function that takes the state and returns a result value together with a modified state
  -> Eff es a

-- | Apply a function to the result of 'get'.
gets :: State s :> es => (s -> t) -> Eff es t
gets :: (s -> t) -> Eff es t
gets = ((s -> t) -> Eff es s -> Eff es t
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Eff es s
forall s (es :: [(Type -> Type) -> Type -> Type]).
(State s :> es) =>
Eff es s
get)

-- | Modify the value of the state via a function.
modify :: State s :> es => (s -> s) -> Eff es ()
modify :: (s -> s) -> Eff es ()
modify s -> s
f = (s -> ((), s)) -> Eff es ()
forall s (es :: [(Type -> Type) -> Type -> Type]) a.
(State s :> es) =>
(s -> (a, s)) -> Eff es a
state (((), ) (s -> ((), s)) -> (s -> s) -> s -> ((), s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> s
f)

-- * Interpretations

handleIORef :: IOE :> es => IORef s -> Handler (State s) es
handleIORef :: IORef s -> Handler (State s) es
handleIORef IORef s
rs = \case
  State s (Eff esSend) a
Get     -> IORef s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => IORef a -> m a
readIORef IORef s
rs
  Put s
s'  -> IORef s -> s -> Eff es ()
forall (m :: Type -> Type) a. MonadIO m => IORef a -> a -> m ()
writeIORef IORef s
rs s
s'
  State s -> (a, s)
f -> IO a -> Eff es a
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO a -> Eff es a) -> IO a -> Eff es a
forall a b. (a -> b) -> a -> b
$ IORef s -> (s -> (s, a)) -> IO a
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORefCAS IORef s
rs ((a, s) -> (s, a)
forall a b. (a, b) -> (b, a)
swap ((a, s) -> (s, a)) -> (s -> (a, s)) -> s -> (s, a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> (a, s)
f)
{-# INLINE handleIORef #-}

-- | Run the 'State' effect.
--
-- === Caveats
--
-- The 'runState' interpreter is implemented with 'Data.IORef.IORef's and there is no way to do arbitrary
-- atomic transactions. The 'state' operation is atomic though and it is implemented with 'atomicModifyIORefCAS', which
-- can be faster than @atomicModifyIORef@ in contention. For any more complicated cases of atomicity, please build your
-- own effect that uses either @MVar@s or @TVar@s based on your need.
--
-- Unlike @mtl@, in @cleff@ the state /will not revert/ when an error is thrown.
--
-- 'runState' will stop taking care of state operations done on forked threads as soon as the main thread finishes its
-- computation. Any state operation done /before main thread finishes/ is still taken into account.
runState :: s -> Eff (State s ': es) a -> Eff es (a, s)
runState :: s -> Eff (State s : es) a -> Eff es (a, s)
runState s
s Eff (State s : es) a
m = Eff (IOE : es) (a, s) -> Eff es (a, s)
forall (es :: [(Type -> Type) -> Type -> Type]).
Eff (IOE : es) ~> Eff es
thisIsPureTrustMe do
  IORef s
rs <- s -> Eff (IOE : es) (IORef s)
forall (m :: Type -> Type) a. MonadIO m => a -> m (IORef a)
newIORef s
s
  a
x <- Handler (State s) (IOE : es)
-> Eff (State s : es) a -> Eff (IOE : es) a
forall (e' :: (Type -> Type) -> Type -> Type)
       (e :: (Type -> Type) -> Type -> Type)
       (es :: [(Type -> Type) -> Type -> Type]).
Handler e (e' : es) -> Eff (e : es) ~> Eff (e' : es)
reinterpret (IORef s -> Handler (State s) (IOE : es)
forall (es :: [(Type -> Type) -> Type -> Type]) s.
(IOE :> es) =>
IORef s -> Handler (State s) es
handleIORef IORef s
rs) Eff (State s : es) a
m
  s
s' <- IORef s -> Eff (IOE : es) s
forall (m :: Type -> Type) a. MonadIO m => IORef a -> m a
readIORef IORef s
rs
  (a, s) -> Eff (IOE : es) (a, s)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (a
x, s
s')
{-# INLINE runState #-}

-- | Run the 'State' effect in terms of operations on a supplied 'IORef'. The 'state' operation is atomic.
--
-- @since 0.2.1.0
runStateIORef :: IOE :> es => IORef s -> Eff (State s ': es) a -> Eff es a
runStateIORef :: IORef s -> Eff (State s : es) a -> Eff es a
runStateIORef IORef s
rs = Handler (State s) es -> Eff (State s : es) ~> Eff es
forall (e :: (Type -> Type) -> Type -> Type)
       (es :: [(Type -> Type) -> Type -> Type]).
Handler e es -> Eff (e : es) ~> Eff es
interpret (Handler (State s) es -> Eff (State s : es) ~> Eff es)
-> Handler (State s) es -> Eff (State s : es) ~> Eff es
forall a b. (a -> b) -> a -> b
$ IORef s -> Handler (State s) es
forall (es :: [(Type -> Type) -> Type -> Type]) s.
(IOE :> es) =>
IORef s -> Handler (State s) es
handleIORef IORef s
rs
{-# INLINE runStateIORef #-}

-- | Run the 'State' effect in terms of operations on a supplied 'MVar'.
--
-- @since 0.2.1.0
runStateMVar :: IOE :> es => MVar s -> Eff (State s ': es) a -> Eff es a
runStateMVar :: MVar s -> Eff (State s : es) a -> Eff es a
runStateMVar MVar s
rs = Handler (State s) es -> Eff (State s : es) ~> Eff es
forall (e :: (Type -> Type) -> Type -> Type)
       (es :: [(Type -> Type) -> Type -> Type]).
Handler e es -> Eff (e : es) ~> Eff es
interpret \case
  State s (Eff esSend) a
Get     -> MVar s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => MVar a -> m a
readMVar MVar s
rs
  Put s'  -> Eff es s -> Eff es ()
forall (f :: Type -> Type) a. Functor f => f a -> f ()
void (Eff es s -> Eff es ()) -> Eff es s -> Eff es ()
forall a b. (a -> b) -> a -> b
$ MVar s -> s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => MVar a -> a -> m a
swapMVar MVar s
rs s
s'
  State f -> MVar s -> (s -> Eff es (s, a)) -> Eff es a
forall (m :: Type -> Type) a b.
MonadUnliftIO m =>
MVar a -> (a -> m (a, b)) -> m b
modifyMVar MVar s
rs \s
s -> let (a
x, !s
s') = s -> (a, s)
f s
s in (s, a) -> Eff es (s, a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (s
s', a
x)
{-# INLINE runStateMVar #-}

-- | Run the 'State' effect in terms of operations on a supplied 'TVar'.
--
-- @since 0.2.1.0
runStateTVar :: IOE :> es => TVar s -> Eff (State s ': es) a -> Eff es a
runStateTVar :: TVar s -> Eff (State s : es) a -> Eff es a
runStateTVar TVar s
rs = Handler (State s) es -> Eff (State s : es) ~> Eff es
forall (e :: (Type -> Type) -> Type -> Type)
       (es :: [(Type -> Type) -> Type -> Type]).
Handler e es -> Eff (e : es) ~> Eff es
interpret \case
  State s (Eff esSend) a
Get -> TVar s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => TVar a -> m a
readTVarIO TVar s
rs
  Put s' -> STM () -> Eff es ()
forall (m :: Type -> Type) a. MonadIO m => STM a -> m a
atomically (STM () -> Eff es ()) -> STM () -> Eff es ()
forall a b. (a -> b) -> a -> b
$ TVar s -> s -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar s
rs s
s'
  State f -> STM a -> Eff es a
forall (m :: Type -> Type) a. MonadIO m => STM a -> m a
atomically do
    s
s <- TVar s -> STM s
forall a. TVar a -> STM a
readTVar TVar s
rs
    let (a
x, !s
s') = s -> (a, s)
f s
s
    TVar s -> s -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar s
rs s
s'
    a -> STM a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
x
{-# INLINE runStateTVar #-}

-- | Run a 'State' effect in terms of a larger 'State' via a 'Lens''.
zoom :: State t :> es => Lens' t s -> Eff (State s ': es) ~> Eff es
zoom :: Lens' t s -> Eff (State s : es) ~> Eff es
zoom Lens' t s
field = Handler (State s) es -> Eff (State s : es) ~> Eff es
forall (e :: (Type -> Type) -> Type -> Type)
       (es :: [(Type -> Type) -> Type -> Type]).
Handler e es -> Eff (e : es) ~> Eff es
interpret \case
  State s (Eff esSend) a
Get     -> (t -> s) -> Eff es s
forall s (es :: [(Type -> Type) -> Type -> Type]) t.
(State s :> es) =>
(s -> t) -> Eff es t
gets (t -> Getting s t s -> s
forall s a. s -> Getting a s a -> a
^. Getting s t s
Lens' t s
field)
  Put s   -> (t -> t) -> Eff es ()
forall s (es :: [(Type -> Type) -> Type -> Type]).
(State s :> es) =>
(s -> s) -> Eff es ()
modify (t -> (t -> t) -> t
forall a b. a -> (a -> b) -> b
& (s -> Identity s) -> t -> Identity t
Lens' t s
field ((s -> Identity s) -> t -> Identity t) -> s -> t -> t
forall s t a b. ASetter s t a b -> b -> s -> t
.~ s
s)
  State f -> (t -> (a, t)) -> Eff es a
forall s (es :: [(Type -> Type) -> Type -> Type]) a.
(State s :> es) =>
(s -> (a, s)) -> Eff es a
state \t
t -> let (a
a, !s
s) = s -> (a, s)
f (t
t t -> Getting s t s -> s
forall s a. s -> Getting a s a -> a
^. Getting s t s
Lens' t s
field) in (a
a, t
t t -> (t -> t) -> t
forall a b. a -> (a -> b) -> b
& (s -> Identity s) -> t -> Identity t
Lens' t s
field ((s -> Identity s) -> t -> Identity t) -> s -> t -> t
forall s t a b. ASetter s t a b -> b -> s -> t
.~ s
s)
{-# INLINE zoom #-}