{-# 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, zoom
  ) where

import           Cleff
import           Cleff.Internal.Base
import           Data.Atomics        (atomicModifyIORefCAS)
import           Data.Tuple          (swap)
import           Lens.Micro          (Lens', (&), (.~), (^.))
import           UnliftIO.IORef      (newIORef, readIORef, writeIORef)

-- * 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@ approach.
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

-- | 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

-- | Run the 'State' effect.
--
-- __Caveat__: 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 (\case
    State s (Eff esSend) a
Get     -> IORef s -> Eff (IOE : es) s
forall (m :: Type -> Type) a. MonadIO m => IORef a -> m a
readIORef IORef s
rs
    Put s'  -> IORef s -> s -> Eff (IOE : es) ()
forall (m :: Type -> Type) a. MonadIO m => IORef a -> a -> m ()
writeIORef IORef s
rs s
s'
    State f -> IO a -> Eff (IOE : es) a
forall (m :: Type -> Type) a. MonadIO m => IO a -> m a
liftIO (IO a -> Eff (IOE : es) a) -> IO a -> Eff (IOE : 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)) 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 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 #-}