{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TemplateHaskell #-}
module Polysemy.State
(
State (..)
, get
, gets
, put
, modify
, modify'
, runState
, evalState
, execState
, runLazyState
, evalLazyState
, execLazyState
, runStateIORef
, stateToIO
, runStateSTRef
, stateToST
, 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
data State s m a where
Get :: State s m s
Put :: s -> State s m ()
makeSem ''State
gets :: forall s a r. Member (State s) r => (s -> a) -> Sem r a
gets :: (s -> a) -> Sem r a
gets s -> a
f = s -> a
f (s -> a) -> Sem r s -> Sem r a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem r s
forall s (r :: [Effect]). MemberWithError (State s) r => Sem r s
get
{-# INLINABLE gets #-}
modify :: Member (State s) r => (s -> s) -> Sem r ()
modify :: (s -> s) -> Sem r ()
modify s -> s
f = do
s
s <- Sem r s
forall s (r :: [Effect]). MemberWithError (State s) r => Sem r s
get
s -> Sem r ()
forall s (r :: [Effect]).
MemberWithError (State s) r =>
s -> Sem r ()
put (s -> Sem r ()) -> s -> Sem r ()
forall a b. (a -> b) -> a -> b
$ s -> s
f s
s
{-# INLINABLE modify #-}
modify' :: Member (State s) r => (s -> s) -> Sem r ()
modify' :: (s -> s) -> Sem r ()
modify' s -> s
f = do
s
s <- Sem r s
forall s (r :: [Effect]). MemberWithError (State s) r => Sem r s
get
s -> Sem r ()
forall s (r :: [Effect]).
MemberWithError (State s) r =>
s -> Sem r ()
put (s -> Sem r ()) -> s -> Sem r ()
forall a b. (a -> b) -> a -> b
$! s -> s
f s
s
{-# INLINABLE modify' #-}
runState :: s -> Sem (State s ': r) a -> Sem r (s, a)
runState :: s -> Sem (State s : r) a -> Sem r (s, a)
runState = (forall x (m :: * -> *). State s m x -> s -> Sem r (s, x))
-> s -> Sem (State s : r) a -> Sem r (s, a)
forall (e :: Effect) s (r :: [Effect]) a.
(forall x (m :: * -> *). e m x -> s -> Sem r (s, x))
-> s -> Sem (e : r) a -> Sem r (s, a)
stateful ((forall x (m :: * -> *). State s m x -> s -> Sem r (s, x))
-> s -> Sem (State s : r) a -> Sem r (s, a))
-> (forall x (m :: * -> *). State s m x -> s -> Sem r (s, x))
-> s
-> Sem (State s : r) a
-> Sem r (s, a)
forall a b. (a -> b) -> a -> b
$ \case
State s m x
Get -> \s
s -> (s, s) -> Sem r (s, s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, s
s)
Put s -> Sem r (s, ()) -> s -> Sem r (s, ())
forall a b. a -> b -> a
const (Sem r (s, ()) -> s -> Sem r (s, ()))
-> Sem r (s, ()) -> s -> Sem r (s, ())
forall a b. (a -> b) -> a -> b
$ (s, ()) -> Sem r (s, ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, ())
{-# INLINE[3] runState #-}
evalState :: s -> Sem (State s ': r) a -> Sem r a
evalState :: s -> Sem (State s : r) a -> Sem r a
evalState s
s = ((s, a) -> a) -> Sem r (s, a) -> Sem r a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (s, a) -> a
forall a b. (a, b) -> b
snd (Sem r (s, a) -> Sem r a)
-> (Sem (State s : r) a -> Sem r (s, a))
-> Sem (State s : r) a
-> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Sem (State s : r) a -> Sem r (s, a)
forall s (r :: [Effect]) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runState s
s
{-# INLINE evalState #-}
execState :: s -> Sem (State s ': r) a -> Sem r s
execState :: s -> Sem (State s : r) a -> Sem r s
execState s
s = ((s, a) -> s) -> Sem r (s, a) -> Sem r s
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (s, a) -> s
forall a b. (a, b) -> a
fst (Sem r (s, a) -> Sem r s)
-> (Sem (State s : r) a -> Sem r (s, a))
-> Sem (State s : r) a
-> Sem r s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Sem (State s : r) a -> Sem r (s, a)
forall s (r :: [Effect]) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runState s
s
{-# INLINE execState #-}
runLazyState :: s -> Sem (State s ': r) a -> Sem r (s, a)
runLazyState :: s -> Sem (State s : r) a -> Sem r (s, a)
runLazyState = (forall x (m :: * -> *). State s m x -> s -> Sem r (s, x))
-> s -> Sem (State s : r) a -> Sem r (s, a)
forall (e :: Effect) s (r :: [Effect]) a.
(forall x (m :: * -> *). e m x -> s -> Sem r (s, x))
-> s -> Sem (e : r) a -> Sem r (s, a)
lazilyStateful ((forall x (m :: * -> *). State s m x -> s -> Sem r (s, x))
-> s -> Sem (State s : r) a -> Sem r (s, a))
-> (forall x (m :: * -> *). State s m x -> s -> Sem r (s, x))
-> s
-> Sem (State s : r) a
-> Sem r (s, a)
forall a b. (a -> b) -> a -> b
$ \case
State s m x
Get -> \s
s -> (s, s) -> Sem r (s, s)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, s
s)
Put s -> Sem r (s, ()) -> s -> Sem r (s, ())
forall a b. a -> b -> a
const (Sem r (s, ()) -> s -> Sem r (s, ()))
-> Sem r (s, ()) -> s -> Sem r (s, ())
forall a b. (a -> b) -> a -> b
$ (s, ()) -> Sem r (s, ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (s
s, ())
{-# INLINE[3] runLazyState #-}
evalLazyState :: s -> Sem (State s ': r) a -> Sem r a
evalLazyState :: s -> Sem (State s : r) a -> Sem r a
evalLazyState s
s = ((s, a) -> a) -> Sem r (s, a) -> Sem r a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (s, a) -> a
forall a b. (a, b) -> b
snd (Sem r (s, a) -> Sem r a)
-> (Sem (State s : r) a -> Sem r (s, a))
-> Sem (State s : r) a
-> Sem r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Sem (State s : r) a -> Sem r (s, a)
forall s (r :: [Effect]) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runLazyState s
s
{-# INLINE evalLazyState #-}
execLazyState :: s -> Sem (State s ': r) a -> Sem r s
execLazyState :: s -> Sem (State s : r) a -> Sem r s
execLazyState s
s = ((s, a) -> s) -> Sem r (s, a) -> Sem r s
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (s, a) -> s
forall a b. (a, b) -> a
fst (Sem r (s, a) -> Sem r s)
-> (Sem (State s : r) a -> Sem r (s, a))
-> Sem (State s : r) a
-> Sem r s
forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Sem (State s : r) a -> Sem r (s, a)
forall s (r :: [Effect]) a.
s -> Sem (State s : r) a -> Sem r (s, a)
runLazyState s
s
{-# INLINE execLazyState #-}
runStateIORef
:: forall s r a
. Member (Embed IO) r
=> IORef s
-> Sem (State s ': r) a
-> Sem r a
runStateIORef :: IORef s -> Sem (State s : r) a -> Sem r a
runStateIORef IORef s
ref = (forall x (rInitial :: [Effect]).
State s (Sem rInitial) x -> Sem r x)
-> Sem (State s : r) a -> Sem r a
forall (e :: Effect) (r :: [Effect]) a.
FirstOrder e "interpret" =>
(forall x (rInitial :: [Effect]). e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall x (rInitial :: [Effect]).
State s (Sem rInitial) x -> Sem r x)
-> Sem (State s : r) a -> Sem r a)
-> (forall x (rInitial :: [Effect]).
State s (Sem rInitial) x -> Sem r x)
-> Sem (State s : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
State s (Sem rInitial) x
Get -> IO s -> Sem r s
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO s -> Sem r s) -> IO s -> Sem r s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
ref
Put s -> IO () -> Sem r ()
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO () -> Sem r ()) -> IO () -> Sem r ()
forall a b. (a -> b) -> a -> b
$ IORef s -> s -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef s
ref s
s
{-# INLINE runStateIORef #-}
stateToIO
:: forall s r a
. Member (Embed IO) r
=> s
-> Sem (State s ': r) a
-> Sem r (s, a)
stateToIO :: s -> Sem (State s : r) a -> Sem r (s, a)
stateToIO s
s Sem (State s : r) a
sem = do
IORef s
ref <- IO (IORef s) -> Sem r (IORef s)
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO (IORef s) -> Sem r (IORef s))
-> IO (IORef s) -> Sem r (IORef s)
forall a b. (a -> b) -> a -> b
$ s -> IO (IORef s)
forall a. a -> IO (IORef a)
newIORef s
s
a
res <- IORef s -> Sem (State s : r) a -> Sem r a
forall s (r :: [Effect]) 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 <- IO s -> Sem r s
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO s -> Sem r s) -> IO s -> Sem r s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
ref
(s, a) -> Sem r (s, a)
forall (m :: * -> *) a. Monad m => a -> m a
return (s
end, a
res)
{-# INLINE stateToIO #-}
runStateSTRef
:: forall s st r a
. Member (Embed (ST st)) r
=> STRef st s
-> Sem (State s ': r) a
-> Sem r a
runStateSTRef :: STRef st s -> Sem (State s : r) a -> Sem r a
runStateSTRef STRef st s
ref = (forall x (rInitial :: [Effect]).
State s (Sem rInitial) x -> Sem r x)
-> Sem (State s : r) a -> Sem r a
forall (e :: Effect) (r :: [Effect]) a.
FirstOrder e "interpret" =>
(forall x (rInitial :: [Effect]). e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall x (rInitial :: [Effect]).
State s (Sem rInitial) x -> Sem r x)
-> Sem (State s : r) a -> Sem r a)
-> (forall x (rInitial :: [Effect]).
State s (Sem rInitial) x -> Sem r x)
-> Sem (State s : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
State s (Sem rInitial) x
Get -> ST st s -> Sem r s
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (ST st s -> Sem r s) -> ST st s -> Sem r s
forall a b. (a -> b) -> a -> b
$ STRef st s -> ST st s
forall s a. STRef s a -> ST s a
readSTRef STRef st s
ref
Put s -> ST st () -> Sem r ()
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (ST st () -> Sem r ()) -> ST st () -> Sem r ()
forall a b. (a -> b) -> a -> b
$ STRef st s -> s -> ST st ()
forall s a. STRef s a -> a -> ST s ()
writeSTRef STRef st s
ref s
s
{-# INLINE runStateSTRef #-}
stateToST
:: forall s st r a
. Member (Embed (ST st)) r
=> s
-> Sem (State s ': r) a
-> Sem r (s, a)
stateToST :: 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 (r :: [Effect]) a.
Member (Embed (ST st)) r =>
ST st a -> Sem r a
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed @(ST st) (ST st (STRef st s) -> Sem r (STRef st s))
-> ST st (STRef st s) -> Sem r (STRef st s)
forall a b. (a -> b) -> a -> b
$ s -> ST st (STRef st s)
forall a s. a -> ST s (STRef s a)
newSTRef s
s
a
res <- STRef st s -> Sem (State s : r) a -> Sem r a
forall s st (r :: [Effect]) 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 <- ST st s -> Sem r s
forall (m :: * -> *) (r :: [Effect]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (ST st s -> Sem r s) -> ST st s -> Sem r s
forall a b. (a -> b) -> a -> b
$ STRef st s -> ST st s
forall s a. STRef s a -> ST s a
readSTRef STRef st s
ref
(s, a) -> Sem r (s, a)
forall (m :: * -> *) a. Monad m => a -> m a
return (s
end, a
res)
{-# INLINE stateToST #-}
hoistStateIntoStateT
:: Sem (State s ': r) a
-> S.StateT s (Sem r) a
hoistStateIntoStateT :: 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 x.
Union (State s : r) (Sem (State s : r)) x -> StateT s (Sem r) x)
-> StateT s (Sem r) a
forall (m :: * -> *).
Monad m =>
(forall x. Union (State s : r) (Sem (State s : r)) x -> m x) -> m a
m ((forall x.
Union (State s : r) (Sem (State s : r)) x -> StateT s (Sem r) x)
-> StateT s (Sem r) a)
-> (forall x.
Union (State s : r) (Sem (State s : r)) x -> StateT s (Sem r) x)
-> StateT s (Sem r) a
forall a b. (a -> b) -> a -> b
$ \Union (State s : r) (Sem (State s : r)) x
u ->
case Union (State s : r) (Sem (State s : r)) x
-> Either
(Union r (Sem (State s : r)) x)
(Weaving (State s) (Sem (State s : r)) x)
forall (e :: Effect) (r :: [Effect]) (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 -> (s -> Sem r (x, s)) -> StateT s (Sem r) x
forall s (m :: * -> *) a. (s -> m (a, s)) -> StateT s m a
S.StateT ((s -> Sem r (x, s)) -> StateT s (Sem r) x)
-> (s -> Sem r (x, s)) -> StateT s (Sem r) x
forall a b. (a -> b) -> a -> b
$ \s
s ->
Union r (Sem r) (x, s) -> Sem r (x, s)
forall (r :: [Effect]) a. Union r (Sem r) a -> Sem r a
liftSem (Union r (Sem r) (x, s) -> Sem r (x, s))
-> (Union r (StateT s (Sem r)) x -> Union r (Sem r) (x, s))
-> Union r (StateT s (Sem r)) x
-> Sem r (x, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((s, x) -> (x, s))
-> Union r (Sem r) (s, x) -> Union r (Sem r) (x, s)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (s, x) -> (x, s)
forall a b. (a, b) -> (b, a)
swap
(Union r (Sem r) (s, x) -> Union r (Sem r) (x, s))
-> (Union r (StateT s (Sem r)) x -> Union r (Sem r) (s, x))
-> Union r (StateT s (Sem r)) x
-> Union r (Sem r) (x, s)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (s, ())
-> (forall x. (s, StateT s (Sem r) x) -> Sem r (s, x))
-> (forall x. (s, x) -> Maybe x)
-> Union r (StateT s (Sem r)) x
-> Union r (Sem r) (s, x)
forall (s :: * -> *) (n :: * -> *) (m :: * -> *) (r :: [Effect]) 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') -> (x, s) -> (s, x)
forall a b. (a, b) -> (b, a)
swap ((x, s) -> (s, x)) -> Sem r (x, s) -> Sem r (s, x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT s (Sem r) x -> s -> Sem r (x, s)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
S.runStateT StateT s (Sem r) x
m' s
s')
(x -> Maybe x
forall a. a -> Maybe a
Just (x -> Maybe x) -> ((s, x) -> x) -> (s, x) -> Maybe x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (s, x) -> x
forall a b. (a, b) -> b
snd)
(Union r (StateT s (Sem r)) x -> Sem r (x, s))
-> Union r (StateT s (Sem r)) x -> Sem r (x, s)
forall a b. (a -> b) -> a -> b
$ (forall x. Sem (State s : r) x -> StateT s (Sem r) x)
-> Union r (Sem (State s : r)) x -> Union r (StateT s (Sem r)) x
forall (m :: * -> *) (n :: * -> *) (r :: [Effect]) a.
(forall x. m x -> n x) -> Union r m a -> Union r n a
hoist forall x. Sem (State s : r) x -> StateT s (Sem r) x
forall s (r :: [Effect]) 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 (f a -> x) -> (a -> f a) -> a -> x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f () -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f ()
z) (a -> x) -> StateT a (Sem r) a -> StateT a (Sem r) x
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT a (Sem r) a
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 (f a -> x) -> (a -> f a) -> a -> x
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f () -> f a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ f ()
z) (a -> x) -> StateT s (Sem r) a -> StateT s (Sem r) x
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> StateT s (Sem r) ()
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
#-}