{-# LANGUAGE TemplateHaskell #-}
module Control.Effect.State
(
State'(..)
, State
, get
, put
, state
, gets'
, gets
, modify'
, modify
, modifyStrict'
, modifyStrict
, tagState'
, retagState'
, untagState'
) where
import Data.Tuple (swap)
import qualified Control.Monad.Trans.RWS.CPS as Strict
import qualified Control.Monad.Trans.RWS.Lazy as Lazy
import qualified Control.Monad.Trans.State.Lazy as L
import qualified Control.Monad.Trans.State.Strict as S
import Control.Effect.Machinery
class Monad m => State' tag s m | tag m -> s where
{-# MINIMAL get', put' | state' #-}
get' :: m s
get' = state' @tag (\s -> (s, s))
{-# INLINE get' #-}
put' :: s -> m ()
put' s = state' @tag (\_ -> (s, ()))
{-# INLINE put' #-}
state' :: (s -> (s, a)) -> m a
state' f = do
s <- get' @tag
let ~(s', a) = f s
put' @tag s'
pure a
{-# INLINE state' #-}
makeTaggedEffect ''State'
instance Monad m => State' tag s (L.StateT s m) where
get' = L.get
{-# INLINE get' #-}
put' = L.put
{-# INLINE put' #-}
state' = L.state . fmap swap
{-# INLINE state' #-}
instance Monad m => State' tag s (S.StateT s m) where
get' = S.get
{-# INLINE get' #-}
put' = S.put
{-# INLINE put' #-}
state' = S.state . fmap swap
{-# INLINE state' #-}
instance (Monad m, Monoid w) => State' tag s (Lazy.RWST r w s m) where
get' = Lazy.get
{-# INLINE get' #-}
put' = Lazy.put
{-# INLINE put' #-}
state' = Lazy.state . fmap swap
{-# INLINE state' #-}
instance Monad m => State' tag s (Strict.RWST r w s m) where
get' = Strict.get
{-# INLINE get' #-}
put' = Strict.put
{-# INLINE put' #-}
state' = Strict.state . fmap swap
{-# INLINE state' #-}
gets' :: forall tag s m a. State' tag s m => (s -> a) -> m a
gets' f = fmap f (get' @tag)
{-# INLINE gets' #-}
gets :: State s m => (s -> a) -> m a
gets = gets' @G
{-# INLINE gets #-}
modify' :: forall tag s m. State' tag s m => (s -> s) -> m ()
modify' f = do
s <- get' @tag
put' @tag (f s)
{-# INLINE modify' #-}
modify :: State s m => (s -> s) -> m ()
modify = modify' @G
{-# INLINE modify #-}
modifyStrict' :: forall tag s m. State' tag s m => (s -> s) -> m ()
modifyStrict' f = do
s <- get' @tag
put' @tag $! f s
{-# INLINE modifyStrict' #-}
modifyStrict :: State s m => (s -> s) -> m ()
modifyStrict = modifyStrict' @G
{-# INLINE modifyStrict #-}