{-# 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
  , runStateLocal
  , runStateIORef
  , runStateMVar
  , runStateTVar
  , zoom
  ) where

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

-- * 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 :: [Effect]). (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 :: [Effect]) 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     -> IO s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO s -> Eff es s) -> IO s -> Eff es s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
rs
  Put s
s'  -> IO () -> Eff es ()
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO () -> Eff es ()) -> IO () -> Eff es ()
forall a b. (a -> b) -> a -> b
$ IORef s -> s -> IO ()
forall a. IORef a -> a -> IO ()
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)

-- | 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 :: [Effect]). Eff (IOE : es) ~> Eff es
thisIsPureTrustMe do
  IORef s
rs <- IO (IORef s) -> Eff (IOE : es) (IORef s)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (IORef s) -> Eff (IOE : es) (IORef s))
-> IO (IORef s) -> Eff (IOE : es) (IORef s)
forall a b. (a -> b) -> a -> b
$ s -> IO (IORef s)
forall a. a -> IO (IORef a)
newIORef s
s
  a
x <- Handler (State s) (IOE : es)
-> Eff (State s : es) a -> Eff (IOE : es) a
forall (e' :: Effect) (e :: Effect) (es :: [Effect]).
Handler e (e' : es) -> Eff (e : es) ~> Eff (e' : es)
reinterpret (IORef s -> Handler (State s) (IOE : es)
forall (es :: [Effect]) s.
(IOE :> es) =>
IORef s -> Handler (State s) es
handleIORef IORef s
rs) Eff (State s : es) a
m
  s
s' <- IO s -> Eff (IOE : es) s
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO s -> Eff (IOE : es) s) -> IO s -> Eff (IOE : es) s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO 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')

-- | Run a 'State' effect where each thread has its thread-local state.
--
-- This means that each thread will have an individual state that has the same initial value. Threfore, state
-- operations on one thread will not change the state for any other thread.
--
-- The returned final state is that of the current thread.
--
-- === Caveats
--
-- Like 'runState', the 'state' operation in this handler is atomic. Like 'runState', and unlike @mtl@, any errors will
-- not revert the state changes.
--
-- Be warned that if you use a thread pool, then when a thread is reused, it may read the state left from the last
-- usage, therefore losing locality. If you use a thread pool, you will want to manually reset the state after each
-- task.
--
-- @since 0.3.3.0
runStateLocal :: s -> Eff (State s : es) a -> Eff es (a, s)
runStateLocal :: s -> Eff (State s : es) a -> Eff es (a, s)
runStateLocal s
s Eff (State s : es) a
m = Eff (IOE : es) (a, s) -> Eff es (a, s)
forall (es :: [Effect]). Eff (IOE : es) ~> Eff es
thisIsPureTrustMe do
  ThreadVar s
rs <- IO (ThreadVar s) -> Eff (IOE : es) (ThreadVar s)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO (ThreadVar s) -> Eff (IOE : es) (ThreadVar s))
-> IO (ThreadVar s) -> Eff (IOE : es) (ThreadVar s)
forall a b. (a -> b) -> a -> b
$ s -> IO (ThreadVar s)
forall a. a -> IO (ThreadVar a)
newThreadVar s
s
  a
x <- Handler (State s) (IOE : es)
-> Eff (State s : es) a -> Eff (IOE : es) a
forall (e' :: Effect) (e :: Effect) (es :: [Effect]).
Handler e (e' : es) -> Eff (e : es) ~> Eff (e' : es)
reinterpret (\State s (Eff esSend) a
e -> IO (IORef s) -> Eff (IOE : es) (IORef s)
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (ThreadVar s -> IO (IORef s)
forall a. ThreadVar a -> IO (IORef a)
getThreadVar ThreadVar s
rs) Eff (IOE : es) (IORef s)
-> (IORef s -> Eff (IOE : es) a) -> Eff (IOE : es) a
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \IORef s
r -> IORef s -> State s (Eff esSend) a -> Eff (IOE : es) a
forall (es :: [Effect]) s.
(IOE :> es) =>
IORef s -> Handler (State s) es
handleIORef IORef s
r State s (Eff esSend) a
e) Eff (State s : es) a
m
  s
s' <- IO s -> Eff (IOE : es) s
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO s -> Eff (IOE : es) s) -> IO s -> Eff (IOE : es) s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO a
readIORef (IORef s -> IO s) -> IO (IORef s) -> IO s
forall (m :: Type -> Type) a b. Monad m => (a -> m b) -> m a -> m b
=<< ThreadVar s -> IO (IORef s)
forall a. ThreadVar a -> IO (IORef a)
getThreadVar ThreadVar s
rs
  (a, s) -> Eff (IOE : es) (a, s)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (a
x, s
s')

-- | 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 :: Effect) (es :: [Effect]).
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 :: [Effect]) s.
(IOE :> es) =>
IORef s -> Handler (State s) es
handleIORef IORef s
rs

-- | 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 :: Effect) (es :: [Effect]).
Handler e es -> Eff (e : es) ~> Eff es
interpret \case
  State s (Eff esSend) a
Get     -> IO s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO s -> Eff es s) -> IO s -> Eff es s
forall a b. (a -> b) -> a -> b
$ MVar s -> IO s
forall a. MVar a -> IO a
readMVar MVar s
rs
  Put s'  -> IO () -> Eff es ()
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO () -> Eff es ()) -> IO () -> Eff es ()
forall a b. (a -> b) -> a -> b
$ IO s -> IO ()
forall (f :: Type -> Type) a. Functor f => f a -> f ()
void (IO s -> IO ()) -> IO s -> IO ()
forall a b. (a -> b) -> a -> b
$ MVar s -> s -> IO s
forall a. MVar a -> a -> IO a
swapMVar MVar s
rs s
s'
  State 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
$ MVar s -> (s -> IO (s, a)) -> IO a
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar s
rs \s
s -> let (a
x, !s
s') = s -> (a, s)
f s
s in (s, a) -> IO (s, a)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (s
s', a
x)

-- | 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 :: Effect) (es :: [Effect]).
Handler e es -> Eff (e : es) ~> Eff es
interpret \case
  State s (Eff esSend) a
Get -> IO s -> Eff es s
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO s -> Eff es s) -> IO s -> Eff es s
forall a b. (a -> b) -> a -> b
$ TVar s -> IO s
forall a. TVar a -> IO a
readTVarIO TVar s
rs
  Put s' -> IO () -> Eff es ()
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO () -> Eff es ()) -> IO () -> Eff es ()
forall a b. (a -> b) -> a -> b
$ STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
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 -> 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
$ STM a -> IO a
forall a. STM a -> IO 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

-- | 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 :: Effect) (es :: [Effect]).
Handler e es -> Eff (e : es) ~> Eff es
interpret \case
  State s (Eff esSend) a
Get     -> (t -> s) -> Eff es s
forall s (es :: [Effect]) 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 :: [Effect]). (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 :: [Effect]) 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)