{-# LANGUAGE CPP #-}
module Control.Effect.AtomicState
  ( -- * Effects
    AtomicState(..)

    -- * Actions
  , atomicState
  , atomicState'
  , atomicGet
  , atomicGets
  , atomicModify
  , atomicModify'

  , atomicPut

    -- * Interpretations
  , atomicStateToIO
  , runAtomicStateIORef
  , runAtomicStateTVar

  , atomicStateToState

    -- * Simple variants of interpretations
  , atomicStateToIOSimple
  , runAtomicStateIORefSimple
  , runAtomicStateTVarSimple

    -- * Carriers
  , AtomicStateToStateC
  ) where

import Data.IORef
import Control.Concurrent.STM

import Control.Effect
import Control.Effect.State

#if MIN_VERSION_base(4,13,0)
import GHC.IORef (atomicModifyIORefP)
#else

data Box a = Box a

atomicModifyIORefP :: IORef s -> (s -> (s, a)) -> IO a
atomicModifyIORefP ref f = do
  Box a <- atomicModifyIORef ref $ \s -> let !(s', a) = f s in (s', Box a)
  return a
{-# INLINE atomicModifyIORefP #-}
#endif

-- | An effect for atomically reading and modifying a piece of state.
--
-- Convention: the interpreter for the @AtomicState@ action must force
-- the resulting tuple of the function, but not the end state or returned value.
data AtomicState s :: Effect where
  AtomicState :: (s -> (s, a)) -> AtomicState s m a
  AtomicGet   :: AtomicState s m s

-- | Atomically read and modify the state.
--
-- The resulting tuple of the computation is forced. You can
-- control what parts of the computation are evaluated by tying
-- their evaluation to the tuple.
atomicState :: Eff (AtomicState s) m => (s -> (s, a)) -> m a
atomicState :: (s -> (s, a)) -> m a
atomicState = AtomicState s m a -> m a
forall (e :: Effect) (m :: * -> *) a.
(Member e (Derivs m), Carrier m) =>
e m a -> m a
send (AtomicState s m a -> m a)
-> ((s -> (s, a)) -> AtomicState s m a) -> (s -> (s, a)) -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (s -> (s, a)) -> AtomicState s m a
forall s a (m :: * -> *). (s -> (s, a)) -> AtomicState s m a
AtomicState
{-# INLINE atomicState #-}

-- | Atomically read and strictly modify the state.
--
-- The resulting state -- but not the value returned -- is forced.
atomicState' :: Eff (AtomicState s) m => (s -> (s, a)) -> m a
atomicState' :: (s -> (s, a)) -> m a
atomicState' s -> (s, a)
f = (s -> (s, a)) -> m a
forall s (m :: * -> *) a.
Eff (AtomicState s) m =>
(s -> (s, a)) -> m a
atomicState ((s -> (s, a)) -> m a) -> (s -> (s, a)) -> m a
forall a b. (a -> b) -> a -> b
$ \s
s -> let (!s
s', a
a) = s -> (s, a)
f s
s in (s
s', a
a)
{-# INLINE atomicState' #-}

-- | Read the state.
--
-- Depending on the interperation of 'AtomicState', this
-- can be more efficient than @'atomicState' (\\s -> (s,s))@
atomicGet :: Eff (AtomicState s) m => m s
atomicGet :: m s
atomicGet = AtomicState s m s -> m s
forall (e :: Effect) (m :: * -> *) a.
(Member e (Derivs m), Carrier m) =>
e m a -> m a
send AtomicState s m s
forall s (m :: * -> *). AtomicState s m s
AtomicGet
{-# INLINE atomicGet #-}

atomicGets :: Eff (AtomicState s) m => (s -> a) -> m a
atomicGets :: (s -> a) -> m a
atomicGets = ((s -> a) -> m s -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m s
forall s (m :: * -> *). Eff (AtomicState s) m => m s
atomicGet)
{-# INLINE atomicGets #-}

-- | Atomically modify the state.
--
-- The resulting state is not forced. 'atomicModify''
-- is a strict version that does force it.
atomicModify :: Eff (AtomicState s) m => (s -> s) -> m ()
atomicModify :: (s -> s) -> m ()
atomicModify s -> s
f = (s -> (s, ())) -> m ()
forall s (m :: * -> *) a.
Eff (AtomicState s) m =>
(s -> (s, a)) -> m a
atomicState ((s -> (s, ())) -> m ()) -> (s -> (s, ())) -> m ()
forall a b. (a -> b) -> a -> b
$ \s
s -> (s -> s
f s
s, ())
{-# INLINE atomicModify #-}

-- | Atomically and strictly modify the state.
--
-- This is a strict version of 'atomicModify'.
atomicModify' :: Eff (AtomicState s) m => (s -> s) -> m ()
atomicModify' :: (s -> s) -> m ()
atomicModify' s -> s
f = (s -> (s, ())) -> m ()
forall s (m :: * -> *) a.
Eff (AtomicState s) m =>
(s -> (s, a)) -> m a
atomicState ((s -> (s, ())) -> m ()) -> (s -> (s, ())) -> m ()
forall a b. (a -> b) -> a -> b
$ \s
s -> let !s' :: s
s' = s -> s
f s
s in (s
s', ())
{-# INLINE atomicModify' #-}

-- | Atomically overwrite the state.
--
-- You typically don't want to use this, as
-- @'atomicGet' >>= 'atomicPut' . f@ isn't atomic.
atomicPut :: Eff (AtomicState s) m => s -> m ()
atomicPut :: s -> m ()
atomicPut s
s = (s -> (s, ())) -> m ()
forall s (m :: * -> *) a.
Eff (AtomicState s) m =>
(s -> (s, a)) -> m a
atomicState ((s -> (s, ())) -> m ()) -> (s -> (s, ())) -> m ()
forall a b. (a -> b) -> a -> b
$ \s
_ -> (s
s, ())
{-# INLINE atomicPut #-}

-- | Run an 'AtomicState' effect in terms of atomic operations in IO.
--
-- Internally, this simply creates a new 'IORef', passes it to
-- 'runAtomicStateIORef', and then returns the result and the final value
-- of the 'IORef'.
--
-- This has a higher-rank type, as it makes use of 'InterpretReifiedC'.
-- __This makes 'atomicStateToIO' very difficult to use partially applied.__
-- __In particular, it can't be composed using @'.'@.__
--
-- If performance is secondary, consider using the slower
-- 'atomicStateToIOSimple', which doesn't have a higher-rank type.
atomicStateToIO :: forall s m a
                 . Eff (Embed IO) m
                => s
                -> InterpretReifiedC (AtomicState s) m a
                -> m (s, a)
atomicStateToIO :: s -> InterpretReifiedC (AtomicState s) m a -> m (s, a)
atomicStateToIO s
sInit InterpretReifiedC (AtomicState s) m a
main = do
  IORef s
ref  <- IO (IORef s) -> m (IORef s)
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IO (IORef s) -> m (IORef s)) -> IO (IORef s) -> m (IORef s)
forall a b. (a -> b) -> a -> b
$ s -> IO (IORef s)
forall a. a -> IO (IORef a)
newIORef s
sInit
  a
a    <- IORef s -> InterpretReifiedC (AtomicState s) m a -> m a
forall s (m :: * -> *) a.
Eff (Embed IO) m =>
IORef s -> InterpretReifiedC (AtomicState s) m a -> m a
runAtomicStateIORef IORef s
ref InterpretReifiedC (AtomicState s) m a
main
  s
sEnd <- IO s -> m s
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IO s -> m s) -> IO s -> m s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
ref
  (s, a) -> m (s, a)
forall (m :: * -> *) a. Monad m => a -> m a
return (s
sEnd, a
a)
{-# INLINE atomicStateToIO #-}

-- | Run an 'AtomicState' effect in terms of atomic operations in IO.
--
-- Internally, this simply creates a new 'IORef', passes it to
-- 'runAtomicStateIORefSimple', and then returns the result and the final value
-- of the 'IORef'.
--
-- This is a less performant version of 'runAtomicStateIORefSimple' that doesn't
-- have a higher-rank type, making it much easier to use partially applied.
atomicStateToIOSimple :: forall s m a p
                       . ( Eff (Embed IO) m
                         , Threaders '[ReaderThreads] m p
                         )
                      => s
                      -> InterpretSimpleC (AtomicState s) m a
                      -> m (s, a)
atomicStateToIOSimple :: s -> InterpretSimpleC (AtomicState s) m a -> m (s, a)
atomicStateToIOSimple s
sInit InterpretSimpleC (AtomicState s) m a
main = do
  IORef s
ref  <- IO (IORef s) -> m (IORef s)
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IO (IORef s) -> m (IORef s)) -> IO (IORef s) -> m (IORef s)
forall a b. (a -> b) -> a -> b
$ s -> IO (IORef s)
forall a. a -> IO (IORef a)
newIORef s
sInit
  a
a    <- IORef s -> InterpretSimpleC (AtomicState s) m a -> m a
forall s (m :: * -> *) a (p :: [Effect]).
(Eff (Embed IO) m, Threaders '[ReaderThreads] m p) =>
IORef s -> InterpretSimpleC (AtomicState s) m a -> m a
runAtomicStateIORefSimple IORef s
ref InterpretSimpleC (AtomicState s) m a
main
  s
sEnd <- IO s -> m s
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IO s -> m s) -> IO s -> m s
forall a b. (a -> b) -> a -> b
$ IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
ref
  (s, a) -> m (s, a)
forall (m :: * -> *) a. Monad m => a -> m a
return (s
sEnd, a
a)
{-# INLINE atomicStateToIOSimple #-}

-- | Run an 'AtomicState' effect by transforming it into atomic operations
-- over an 'IORef'.
--
-- This has a higher-rank type, as it makes use of 'InterpretReifiedC'.
-- __This makes 'runAtomicStateIORef' very difficult to use partially applied.__
-- __In particular, it can't be composed using @'.'@.__
--
-- If performance is secondary, consider using the slower
-- 'runAtomicStateIORefSimple', which doesn't have a higher-rank type.
runAtomicStateIORef :: forall s m a
                     . Eff (Embed IO) m
                    => IORef s
                    -> InterpretReifiedC (AtomicState s) m a
                    -> m a
runAtomicStateIORef :: IORef s -> InterpretReifiedC (AtomicState s) m a -> m a
runAtomicStateIORef IORef s
ref = EffHandler (AtomicState s) m
-> InterpretReifiedC (AtomicState s) m a -> m a
forall (e :: Effect) (m :: * -> *) a.
(RepresentationalEff e, Carrier m) =>
EffHandler e m -> InterpretReifiedC e m a -> m a
interpret (EffHandler (AtomicState s) m
 -> InterpretReifiedC (AtomicState s) m a -> m a)
-> EffHandler (AtomicState s) m
-> InterpretReifiedC (AtomicState s) m a
-> m a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> IO x -> Effly z x
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IORef s -> (s -> (s, x)) -> IO x
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORefP IORef s
ref s -> (s, x)
f)
  AtomicState s (Effly z) x
AtomicGet     -> IO s -> Effly z s
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
ref)
{-# INLINE runAtomicStateIORef #-}

-- | Run an 'AtomicState' effect by transforming it into atomic operations
-- over an 'IORef'.
--
-- This is a less performant version of 'runAtomicStateIORef' that doesn't have
-- a higher-rank type, making it much easier to use partially applied.
runAtomicStateIORefSimple :: forall s m a p
                           . ( Eff (Embed IO) m
                             , Threaders '[ReaderThreads] m p
                             )
                          => IORef s
                          -> InterpretSimpleC (AtomicState s) m a
                          -> m a
runAtomicStateIORefSimple :: IORef s -> InterpretSimpleC (AtomicState s) m a -> m a
runAtomicStateIORefSimple IORef s
ref = EffHandler (AtomicState s) m
-> InterpretSimpleC (AtomicState s) m a -> m a
forall (e :: Effect) (m :: * -> *) a (p :: [Effect]).
(RepresentationalEff e, Threaders '[ReaderThreads] m p,
 Carrier m) =>
EffHandler e m -> InterpretSimpleC e m a -> m a
interpretSimple (EffHandler (AtomicState s) m
 -> InterpretSimpleC (AtomicState s) m a -> m a)
-> EffHandler (AtomicState s) m
-> InterpretSimpleC (AtomicState s) m a
-> m a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> IO x -> Effly z x
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IORef s -> (s -> (s, x)) -> IO x
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORefP IORef s
ref s -> (s, x)
f)
  AtomicState s (Effly z) x
AtomicGet     -> IO s -> Effly z s
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IORef s -> IO s
forall a. IORef a -> IO a
readIORef IORef s
ref)
{-# INLINE runAtomicStateIORefSimple #-}

-- | Run an 'AtomicState' effect by transforming it into atomic operations
-- over an 'TVar'.
--
-- This has a higher-rank type, as it makes use of 'InterpretReifiedC'.
-- __This makes 'runAtomicStateTVar' very difficult to use partially applied.__
-- __In particular, it can't be composed using @'.'@.__
--
-- If performance is secondary, consider using the slower
-- 'runAtomicStateTVarSimple', which doesn't have a higher-rank type.
runAtomicStateTVar :: forall s m a
                    . Eff (Embed IO) m
                   => TVar s
                   -> InterpretReifiedC (AtomicState s) m a
                   -> m a
runAtomicStateTVar :: TVar s -> InterpretReifiedC (AtomicState s) m a -> m a
runAtomicStateTVar TVar s
tvar = EffHandler (AtomicState s) m
-> InterpretReifiedC (AtomicState s) m a -> m a
forall (e :: Effect) (m :: * -> *) a.
(RepresentationalEff e, Carrier m) =>
EffHandler e m -> InterpretReifiedC e m a -> m a
interpret (EffHandler (AtomicState s) m
 -> InterpretReifiedC (AtomicState s) m a -> m a)
-> EffHandler (AtomicState s) m
-> InterpretReifiedC (AtomicState s) m a
-> m a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> IO x -> Effly z x
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IO x -> Effly z x) -> IO x -> Effly z x
forall a b. (a -> b) -> a -> b
$ STM x -> IO x
forall a. STM a -> IO a
atomically (STM x -> IO x) -> STM x -> IO x
forall a b. (a -> b) -> a -> b
$ do
    (s
s, x
a) <- s -> (s, x)
f (s -> (s, x)) -> STM s -> STM (s, x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar s -> STM s
forall a. TVar a -> STM a
readTVar TVar s
tvar
    TVar s -> s -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar s
tvar s
s
    x -> STM x
forall (m :: * -> *) a. Monad m => a -> m a
return x
a
  AtomicState s (Effly z) x
AtomicGet     -> IO s -> Effly z s
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (TVar s -> IO s
forall a. TVar a -> IO a
readTVarIO TVar s
tvar)
{-# INLINE runAtomicStateTVar #-}

-- | Run an 'AtomicState' effect by transforming it into atomic operations
-- over an 'TVar'.
--
-- This is a less performant version of 'runAtomicStateIORef' that doesn't have
-- a higher-rank type, making it much easier to use partially applied.
runAtomicStateTVarSimple :: forall s m a p
                          . ( Eff (Embed IO) m
                            , Threaders '[ReaderThreads] m p
                            )
                         => TVar s
                         -> InterpretSimpleC (AtomicState s) m a
                         -> m a
runAtomicStateTVarSimple :: TVar s -> InterpretSimpleC (AtomicState s) m a -> m a
runAtomicStateTVarSimple TVar s
tvar = EffHandler (AtomicState s) m
-> InterpretSimpleC (AtomicState s) m a -> m a
forall (e :: Effect) (m :: * -> *) a (p :: [Effect]).
(RepresentationalEff e, Threaders '[ReaderThreads] m p,
 Carrier m) =>
EffHandler e m -> InterpretSimpleC e m a -> m a
interpretSimple (EffHandler (AtomicState s) m
 -> InterpretSimpleC (AtomicState s) m a -> m a)
-> EffHandler (AtomicState s) m
-> InterpretSimpleC (AtomicState s) m a
-> m a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> IO x -> Effly z x
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (IO x -> Effly z x) -> IO x -> Effly z x
forall a b. (a -> b) -> a -> b
$ STM x -> IO x
forall a. STM a -> IO a
atomically (STM x -> IO x) -> STM x -> IO x
forall a b. (a -> b) -> a -> b
$ do
    (s
s, x
a) <- s -> (s, x)
f (s -> (s, x)) -> STM s -> STM (s, x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar s -> STM s
forall a. TVar a -> STM a
readTVar TVar s
tvar
    TVar s -> s -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar s
tvar s
s
    x -> STM x
forall (m :: * -> *) a. Monad m => a -> m a
return x
a
  AtomicState s (Effly z) x
AtomicGet     -> IO s -> Effly z s
forall (b :: * -> *) (m :: * -> *) a. Eff (Embed b) m => b a -> m a
embed (TVar s -> IO s
forall a. TVar a -> IO a
readTVarIO TVar s
tvar)
{-# INLINE runAtomicStateTVarSimple #-}

data AtomicStateToStateH

type AtomicStateToStateC s = InterpretC AtomicStateToStateH (AtomicState s)

instance Eff (State s) m
      => Handler AtomicStateToStateH (AtomicState s) m where
  effHandler :: AtomicState s (Effly z) x -> Effly z x
effHandler = \case
    AtomicState s -> (s, x)
f -> (s -> (s, x)) -> Effly z x
forall s (m :: * -> *) a. Eff (State s) m => (s -> (s, a)) -> m a
state s -> (s, x)
f
    AtomicState s (Effly z) x
AtomicGet     -> Effly z x
forall s (m :: * -> *). Eff (State s) m => m s
get
  {-# INLINEABLE effHandler #-}

-- | Transform an 'AtomicState' effect into a 'State' effect, discarding atomicity.
atomicStateToState :: Eff (State s) m
                   => AtomicStateToStateC s m a
                   -> m a
atomicStateToState :: AtomicStateToStateC s m a -> m a
atomicStateToState = AtomicStateToStateC s m a -> m a
forall h (e :: Effect) (m :: * -> *) a.
Handler h e m =>
InterpretC h e m a -> m a
interpretViaHandler
{-# INLINE atomicStateToState #-}