{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Description: The 'State' effect
module Polysemy.State
  ( -- * Effect
    State (..)

    -- * Actions
  , get
  , gets
  , put
  , modify
  , modify'

    -- * Interpretations
  , runState
  , evalState
  , execState
  , runLazyState
  , evalLazyState
  , execLazyState
  , runStateIORef
  , stateToIO
  , runStateSTRef
  , stateToST

    -- * Interoperation with MTL
  , hoistStateIntoStateT
  ) where

import Control.Monad.ST
import qualified Control.Monad.Trans.State as S
import Data.IORef
import Data.STRef
import Data.Tuple (swap)
import Polysemy
import Polysemy.Internal
import Polysemy.Internal.Combinators
import Polysemy.Internal.Union


------------------------------------------------------------------------------
-- | An effect for providing statefulness. Note that unlike mtl's
-- 'Control.Monad.Trans.State.StateT', there is no restriction that the 'State'
-- effect corresponds necessarily to /local/ state. It could could just as well
-- be interrpeted in terms of HTTP requests or database access.
--
-- Interpreters which require statefulness can 'Polysemy.reinterpret'
-- themselves in terms of 'State', and subsequently call 'runState'.
data State s m a where
  -- | Get the state.
  Get :: State s m s
  -- | Update the state.
  Put :: s -> State s m ()

makeSem ''State


------------------------------------------------------------------------------
-- | Apply a function to the state and return the result.
gets :: forall s a r. Member (State s) r => (s -> a) -> Sem r a
gets :: forall s a (r :: EffectRow).
Member (State s) r =>
(s -> a) -> Sem r a
gets s -> a
f = s -> a
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (r :: EffectRow). Member (State s) r => Sem r s
get
{-# INLINABLE gets #-}


------------------------------------------------------------------------------
-- | Modify the state.
modify :: Member (State s) r => (s -> s) -> Sem r ()
modify :: forall s (r :: EffectRow).
Member (State s) r =>
(s -> s) -> Sem r ()
modify s -> s
f = do
  s
s <- forall s (r :: EffectRow). Member (State s) r => Sem r s
get
  forall s (r :: EffectRow). Member (State s) r => s -> Sem r ()
put forall a b. (a -> b) -> a -> b
$ s -> s
f s
s
{-# INLINABLE modify #-}

------------------------------------------------------------------------------
-- | A variant of 'modify' in which the computation is strict in the
-- new state.
modify' :: Member (State s) r => (s -> s) -> Sem r ()
modify' :: forall s (r :: EffectRow).
Member (State s) r =>
(s -> s) -> Sem r ()
modify' s -> s
f = do
  s
s <- forall s (r :: EffectRow). Member (State s) r => Sem r s
get
  forall s (r :: EffectRow). Member (State s) r => s -> Sem r ()
put forall a b. (a -> b) -> a -> b
$! s -> s
f s
s
{-# INLINABLE modify' #-}


------------------------------------------------------------------------------
-- | Run a 'State' effect with local state.
runState :: s -> Sem (State s ': r) a -> Sem r (s, a)
runState :: forall s (r :: EffectRow) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runState = forall (e :: Effect) s (r :: EffectRow) a.
(forall x (m :: * -> *). e m x -> s -> Sem r (s, x))
-> s -> Sem (e : r) a -> Sem r (s, a)
stateful forall a b. (a -> b) -> a -> b
$ \case
  State s m x
Get   -> \s
s -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, s
s)
  Put s
s -> forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, ())
{-# INLINE[3] runState #-}


------------------------------------------------------------------------------
-- | Run a 'State' effect with local state.
--
-- @since 1.0.0.0
evalState :: s -> Sem (State s ': r) a -> Sem r a
evalState :: forall s (r :: EffectRow) a. s -> Sem (State s : r) a -> Sem r a
evalState s
s = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (r :: EffectRow) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runState s
s
{-# INLINE evalState #-}

------------------------------------------------------------------------------
-- | Run a 'State' effect with local state.
--
-- @since 1.2.3.1
execState :: s -> Sem (State s ': r) a -> Sem r s
execState :: forall s (r :: EffectRow) a. s -> Sem (State s : r) a -> Sem r s
execState s
s = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (r :: EffectRow) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runState s
s
{-# INLINE execState #-}



------------------------------------------------------------------------------
-- | Run a 'State' effect with local state, lazily.
runLazyState :: s -> Sem (State s ': r) a -> Sem r (s, a)
runLazyState :: forall s (r :: EffectRow) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runLazyState = forall (e :: Effect) s (r :: EffectRow) a.
(forall x (m :: * -> *). e m x -> s -> Sem r (s, x))
-> s -> Sem (e : r) a -> Sem r (s, a)
lazilyStateful forall a b. (a -> b) -> a -> b
$ \case
  State s m x
Get   -> \s
s -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, s
s)
  Put s
s -> forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, ())
{-# INLINE[3] runLazyState #-}

------------------------------------------------------------------------------
-- | Run a 'State' effect with local state, lazily.
--
-- @since 1.0.0.0
evalLazyState :: s -> Sem (State s ': r) a -> Sem r a
evalLazyState :: forall s (r :: EffectRow) a. s -> Sem (State s : r) a -> Sem r a
evalLazyState s
s = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (r :: EffectRow) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runLazyState s
s
{-# INLINE evalLazyState #-}


------------------------------------------------------------------------------
-- | Run a 'State' effect with local state, lazily.
--
-- @since 1.2.3.1
execLazyState :: s -> Sem (State s ': r) a -> Sem r s
execLazyState :: forall s (r :: EffectRow) a. s -> Sem (State s : r) a -> Sem r s
execLazyState s
s = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (r :: EffectRow) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runLazyState s
s
{-# INLINE execLazyState #-}



------------------------------------------------------------------------------
-- | Run a 'State' effect by transforming it into operations over an 'IORef'.
--
-- /Note/: This is not safe in a concurrent setting, as 'modify' isn't atomic.
-- If you need operations over the state to be atomic,
-- use 'Polysemy.AtomicState.runAtomicStateIORef' or
-- 'Polysemy.AtomicState.runAtomicStateTVar' instead.
--
-- @since 1.0.0.0
runStateIORef
    :: forall s r a
     . Member (Embed IO) r
    => IORef s
    -> Sem (State s ': r) a
    -> Sem r a
runStateIORef :: forall s (r :: EffectRow) a.
Member (Embed IO) r =>
IORef s -> Sem (State s : r) a -> Sem r a
runStateIORef IORef s
ref = forall (e :: Effect) (r :: EffectRow) a.
FirstOrder e "interpret" =>
(forall (rInitial :: EffectRow) x. e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret forall a b. (a -> b) -> a -> b
$ \case
  State s (Sem rInitial) x
Get   -> forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef s
ref
  Put s
s -> forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef s
ref s
s
{-# INLINE runStateIORef #-}

--------------------------------------------------------------------
-- | Run an 'State' effect in terms of operations
-- in 'IO'.
--
-- Internally, this simply creates a new 'IORef', passes it to
-- 'runStateIORef', and then returns the result and the final value
-- of the 'IORef'.
--
-- /Note/: This is not safe in a concurrent setting, as 'modify' isn't atomic.
-- If you need operations over the state to be atomic,
-- use 'Polysemy.AtomicState.atomicStateToIO' instead.
--
-- /Beware/: As this uses an 'IORef' internally,
-- all other effects will have local
-- state semantics in regards to 'State' effects
-- interpreted this way.
-- For example, 'Polysemy.Error.throw' and 'Polysemy.Error.catch' will
-- never revert 'put's, even if 'Polysemy.Error.runError' is used
-- after 'stateToIO'.
--
-- @since 1.2.0.0
stateToIO
    :: forall s r a
     . Member (Embed IO) r
    => s
    -> Sem (State s ': r) a
    -> Sem r (s, a)
stateToIO :: forall s (r :: EffectRow) a.
Member (Embed IO) r =>
s -> Sem (State s : r) a -> Sem r (s, a)
stateToIO s
s Sem (State s : r) a
sem = do
  IORef s
ref <- forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef s
s
  a
res <- forall s (r :: EffectRow) a.
Member (Embed IO) r =>
IORef s -> Sem (State s : r) a -> Sem r a
runStateIORef IORef s
ref Sem (State s : r) a
sem
  s
end <- forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef IORef s
ref
  forall (m :: * -> *) a. Monad m => a -> m a
return (s
end, a
res)
{-# INLINE stateToIO #-}

------------------------------------------------------------------------------
-- | Run a 'State' effect by transforming it into operations over an 'STRef'.
--
-- @since 1.3.0.0
runStateSTRef
    :: forall s st r a
     . Member (Embed (ST st)) r
    => STRef st s
    -> Sem (State s ': r) a
    -> Sem r a
runStateSTRef :: forall s st (r :: EffectRow) a.
Member (Embed (ST st)) r =>
STRef st s -> Sem (State s : r) a -> Sem r a
runStateSTRef STRef st s
ref = forall (e :: Effect) (r :: EffectRow) a.
FirstOrder e "interpret" =>
(forall (rInitial :: EffectRow) x. e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret forall a b. (a -> b) -> a -> b
$ \case
  State s (Sem rInitial) x
Get   -> forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall s a. STRef s a -> ST s a
readSTRef STRef st s
ref
  Put s
s -> forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef st s
ref s
s
{-# INLINE runStateSTRef #-}

--------------------------------------------------------------------
-- | Run an 'State' effect in terms of operations
-- in 'ST'.
--
-- Internally, this simply creates a new 'STRef', passes it to
-- 'runStateSTRef', and then returns the result and the final value
-- of the 'STRef'.
--
-- /Beware/: As this uses an 'STRef' internally,
-- all other effects will have local
-- state semantics in regards to 'State' effects
-- interpreted this way.
-- For example, 'Polysemy.Error.throw' and 'Polysemy.Error.catch' will
-- never revert 'put's, even if 'Polysemy.Error.runError' is used
-- after 'stateToST'.
--
-- When not using the plugin, one must introduce the existential @st@ type to
-- 'stateToST', so that the resulting type after 'runM' can be resolved into
-- @forall st. ST st (s, a)@ for use with 'runST'. Doing so requires
-- @-XScopedTypeVariables@.
--
-- @
-- stResult :: forall s a. (s, a)
-- stResult = runST ( (runM $ stateToST \@_ \@st undefined $ pure undefined) :: forall st. ST st (s, a) )
-- @
--
-- @since 1.3.0.0
stateToST
    :: forall s st r a
     . Member (Embed (ST st)) r
    => s
    -> Sem (State s ': r) a
    -> Sem r (s, a)
stateToST :: forall s st (r :: EffectRow) a.
Member (Embed (ST st)) r =>
s -> Sem (State s : r) a -> Sem r (s, a)
stateToST s
s Sem (State s : r) a
sem = do
  STRef st s
ref <- forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed @(ST st) forall a b. (a -> b) -> a -> b
$ forall a s. a -> ST s (STRef s a)
newSTRef s
s
  a
res <- forall s st (r :: EffectRow) a.
Member (Embed (ST st)) r =>
STRef st s -> Sem (State s : r) a -> Sem r a
runStateSTRef STRef st s
ref Sem (State s : r) a
sem
  s
end <- forall (m :: * -> *) (r :: EffectRow) a.
Member (Embed m) r =>
m a -> Sem r a
embed forall a b. (a -> b) -> a -> b
$ forall s a. STRef s a -> ST s a
readSTRef STRef st s
ref
  forall (m :: * -> *) a. Monad m => a -> m a
return (s
end, a
res)
{-# INLINE stateToST #-}

------------------------------------------------------------------------------
-- | Hoist a 'State' effect into a 'S.StateT' monad transformer. This can be
-- useful when writing interpreters that need to interop with MTL.
--
-- @since 0.1.3.0
hoistStateIntoStateT
    :: Sem (State s ': r) a
    -> S.StateT s (Sem r) a
hoistStateIntoStateT :: forall s (r :: EffectRow) a.
Sem (State s : r) a -> StateT s (Sem r) a
hoistStateIntoStateT (Sem forall (m :: * -> *).
Monad m =>
(forall x. Union (State s : r) (Sem (State s : r)) x -> m x) -> m a
m) = forall (m :: * -> *).
Monad m =>
(forall x. Union (State s : r) (Sem (State s : r)) x -> m x) -> m a
m forall a b. (a -> b) -> a -> b
$ \Union (State s : r) (Sem (State s : r)) x
u ->
  case forall (e :: Effect) (r :: EffectRow) (m :: * -> *) a.
Union (e : r) m a -> Either (Union r m a) (Weaving e m a)
decomp Union (State s : r) (Sem (State s : r)) x
u of
    Left Union r (Sem (State s : r)) x
x -> forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
S.StateT forall a b. (a -> b) -> a -> b
$ \s
s ->
      forall (r :: EffectRow) a. Union r (Sem r) a -> Sem r a
liftSem forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> (b, a)
swap
              forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (s :: * -> *) (n :: * -> *) (m :: * -> *) (r :: EffectRow)
       a.
(Functor s, Functor n) =>
s ()
-> (forall x. s (m x) -> n (s x))
-> (forall x. s x -> Maybe x)
-> Union r m a
-> Union r n (s a)
weave (s
s, ())
                      (\(s
s', StateT s (Sem r) x
m') -> forall a b. (a, b) -> (b, a)
swap forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
S.runStateT StateT s (Sem r) x
m' s
s')
                      (forall a. a -> Maybe a
Just forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd)
              forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (n :: * -> *) (r :: EffectRow) a.
(forall x. m x -> n x) -> Union r m a -> Union r n a
hoist forall s (r :: EffectRow) a.
Sem (State s : r) a -> StateT s (Sem r) a
hoistStateIntoStateT Union r (Sem (State s : r)) x
x
    Right (Weaving State s (Sem rInitial) a
Get f ()
z forall x. f (Sem rInitial x) -> Sem (State s : r) (f x)
_ f a -> x
y forall x. f x -> Maybe x
_)     -> f a -> x
y forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f ()
z) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) s. Monad m => StateT s m s
S.get
    Right (Weaving (Put s
s) f ()
z forall x. f (Sem rInitial x) -> Sem (State s : r) (f x)
_ f a -> x
y forall x. f x -> Maybe x
_) -> f a -> x
y forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f ()
z) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) s. Monad m => s -> StateT s m ()
S.put s
s
{-# INLINE hoistStateIntoStateT #-}


{-# RULES "runState/reinterpret"
   forall s e (f :: forall m x. e m x -> Sem (State s ': r) x).
     runState s (reinterpret f e) = stateful (\x s' -> runState s' $ f x) s e
     #-}

{-# RULES "runLazyState/reinterpret"
   forall s e (f :: forall m x. e m x -> Sem (State s ': r) x).
     runLazyState s (reinterpret f e) = lazilyStateful (\x s' -> runLazyState s' $ f x) s e
     #-}