{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TemplateHaskell #-}
module Polysemy.AtomicState
  ( -- * Effect
    AtomicState (..)

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

    -- * Interpretations
  , runAtomicStateIORef
  , runAtomicStateTVar
  , atomicStateToIO
  , atomicStateToState
  ) where


import Control.Concurrent.STM

import Polysemy
import Polysemy.State

import Data.IORef

------------------------------------------------------------------------------
-- | A variant of 'State' that supports atomic operations.
--
-- @since 1.1.0.0
data AtomicState s m a where
  AtomicState :: (s -> (s, a)) -> AtomicState s m a
  AtomicGet   :: AtomicState s m s

makeSem_ ''AtomicState

-----------------------------------------------------------------------------
-- | Atomically reads and modifies the state.
atomicState :: forall s a r
             . Member (AtomicState s) r
            => (s -> (s, a))
            -> Sem r a

atomicGet :: forall s r
           . Member (AtomicState s) r
          => Sem r s

------------------------------------------------------------------------------
-- | @since 1.2.2.0
atomicGets :: forall s s' r
            . Member (AtomicState s) r
           => (s -> s')
           -> Sem r s'
atomicGets :: (s -> s') -> Sem r s'
atomicGets = ((s -> s') -> Sem r s -> Sem r s'
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem r s
forall s (r :: [(* -> *) -> * -> *]).
Member (AtomicState s) r =>
Sem r s
atomicGet)
{-# INLINE atomicGets #-}

-----------------------------------------------------------------------------
-- | A variant of 'atomicState' in which the computation is strict in the new
-- state and return value.
atomicState' :: forall s a r
              . Member (AtomicState s) r
             => (s -> (s, a))
             -> Sem r a
atomicState' :: (s -> (s, a)) -> Sem r a
atomicState' s -> (s, a)
f = do
  -- KingoftheHomeless: return value needs to be forced due to how
  -- 'atomicModifyIORef' is implemented: the computation
  -- (and thus the new state) is forced only once the return value is.
  !a
a <- (s -> (s, a)) -> Sem r a
forall s a (r :: [(* -> *) -> * -> *]).
Member (AtomicState s) r =>
(s -> (s, a)) -> Sem r a
atomicState ((s -> (s, a)) -> Sem r a) -> (s -> (s, a)) -> Sem r a
forall a b. (a -> b) -> a -> b
$ \s
s ->
    case s -> (s, a)
f s
s of
      v :: (s, a)
v@(!s
_, a
_) -> (s, a)
v
  a -> Sem r a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a
{-# INLINE atomicState' #-}

atomicPut :: Member (AtomicState s) r
          => s
          -> Sem r ()
atomicPut :: s -> Sem r ()
atomicPut s
s = do
  !()
_ <- (s -> (s, ())) -> Sem r ()
forall s a (r :: [(* -> *) -> * -> *]).
Member (AtomicState s) r =>
(s -> (s, a)) -> Sem r a
atomicState ((s -> (s, ())) -> Sem r ()) -> (s -> (s, ())) -> Sem r ()
forall a b. (a -> b) -> a -> b
$ \s
_ -> (s
s, ()) -- strict put with atomicModifyIORef
  () -> Sem r ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE atomicPut #-}

atomicModify :: Member (AtomicState s) r
             => (s -> s)
             -> Sem r ()
atomicModify :: (s -> s) -> Sem r ()
atomicModify s -> s
f = (s -> (s, ())) -> Sem r ()
forall s a (r :: [(* -> *) -> * -> *]).
Member (AtomicState s) r =>
(s -> (s, a)) -> Sem r a
atomicState ((s -> (s, ())) -> Sem r ()) -> (s -> (s, ())) -> Sem r ()
forall a b. (a -> b) -> a -> b
$ \s
s -> (s -> s
f s
s, ())
{-# INLINE atomicModify #-}

-----------------------------------------------------------------------------
-- | A variant of 'atomicModify' in which the computation is strict in the
-- new state.
atomicModify' :: Member (AtomicState s) r
              => (s -> s)
              -> Sem r ()
atomicModify' :: (s -> s) -> Sem r ()
atomicModify' s -> s
f = do
  !()
_ <- (s -> (s, ())) -> Sem r ()
forall s a (r :: [(* -> *) -> * -> *]).
Member (AtomicState s) r =>
(s -> (s, a)) -> Sem r a
atomicState ((s -> (s, ())) -> Sem r ()) -> (s -> (s, ())) -> Sem r ()
forall a b. (a -> b) -> a -> b
$ \s
s -> let !s' :: s
s' = s -> s
f s
s in (s
s', ())
  () -> Sem r ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
{-# INLINE atomicModify' #-}

------------------------------------------------------------------------------
-- | Run an 'AtomicState' effect by transforming it into atomic operations
-- over an 'IORef'.
runAtomicStateIORef :: forall s r a
                     . Member (Embed IO) r
                    => IORef s
                    -> Sem (AtomicState s ': r) a
                    -> Sem r a
runAtomicStateIORef :: IORef s -> Sem (AtomicState s : r) a -> Sem r a
runAtomicStateIORef IORef s
ref = (forall x (rInitial :: [(* -> *) -> * -> *]).
 AtomicState s (Sem rInitial) x -> Sem r x)
-> Sem (AtomicState s : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: [(* -> *) -> * -> *]) a.
FirstOrder e "interpret" =>
(forall x (rInitial :: [(* -> *) -> * -> *]).
 e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall x (rInitial :: [(* -> *) -> * -> *]).
  AtomicState s (Sem rInitial) x -> Sem r x)
 -> Sem (AtomicState s : r) a -> Sem r a)
-> (forall x (rInitial :: [(* -> *) -> * -> *]).
    AtomicState s (Sem rInitial) x -> Sem r x)
-> Sem (AtomicState s : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> IO x -> Sem r x
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO x -> Sem r x) -> IO x -> Sem r x
forall a b. (a -> b) -> a -> b
$ IORef s -> (s -> (s, x)) -> IO x
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef s
ref s -> (s, x)
f
  AtomicState s (Sem rInitial) x
AtomicGet     -> IO s -> Sem r s
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) 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
{-# INLINE runAtomicStateIORef #-}

------------------------------------------------------------------------------
-- | Run an 'AtomicState' effect by transforming it into atomic operations
-- over a 'TVar'.
runAtomicStateTVar :: Member (Embed IO) r
                   => TVar s
                   -> Sem (AtomicState s ': r) a
                   -> Sem r a
runAtomicStateTVar :: TVar s -> Sem (AtomicState s : r) a -> Sem r a
runAtomicStateTVar TVar s
tvar = (forall x (rInitial :: [(* -> *) -> * -> *]).
 AtomicState s (Sem rInitial) x -> Sem r x)
-> Sem (AtomicState s : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: [(* -> *) -> * -> *]) a.
FirstOrder e "interpret" =>
(forall x (rInitial :: [(* -> *) -> * -> *]).
 e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall x (rInitial :: [(* -> *) -> * -> *]).
  AtomicState s (Sem rInitial) x -> Sem r x)
 -> Sem (AtomicState s : r) a -> Sem r a)
-> (forall x (rInitial :: [(* -> *) -> * -> *]).
    AtomicState s (Sem rInitial) x -> Sem r x)
-> Sem (AtomicState s : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> IO x -> Sem r x
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) a.
Member (Embed m) r =>
m a -> Sem r a
embed (IO x -> Sem r x) -> IO x -> Sem r 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 (Sem rInitial) x
AtomicGet -> IO s -> Sem r s
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) 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
$ TVar s -> IO s
forall a. TVar a -> IO a
readTVarIO TVar s
tvar
{-# INLINE runAtomicStateTVar #-}

--------------------------------------------------------------------
-- | 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'.
--
-- /Beware/: As this uses an 'IORef' internally,
-- all other effects will have local
-- state semantics in regards to 'AtomicState' effects
-- interpreted this way.
-- For example, 'Polysemy.Error.throw' and 'Polysemy.Error.catch' will
-- never revert 'atomicModify's, even if 'Polysemy.Error.runError' is used
-- after 'atomicStateToIO'.
--
-- @since 1.2.0.0
atomicStateToIO :: forall s r a
                 . Member (Embed IO) r
                => s
                -> Sem (AtomicState s ': r) a
                -> Sem r (s, a)
atomicStateToIO :: s -> Sem (AtomicState s : r) a -> Sem r (s, a)
atomicStateToIO s
s Sem (AtomicState s : r) a
sem = do
  IORef s
ref <- IO (IORef s) -> Sem r (IORef s)
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) 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 (AtomicState s : r) a -> Sem r a
forall s (r :: [(* -> *) -> * -> *]) a.
Member (Embed IO) r =>
IORef s -> Sem (AtomicState s : r) a -> Sem r a
runAtomicStateIORef IORef s
ref Sem (AtomicState s : r) a
sem
  s
end <- IO s -> Sem r s
forall (m :: * -> *) (r :: [(* -> *) -> * -> *]) 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 atomicStateToIO #-}

------------------------------------------------------------------------------
-- | Transform an 'AtomicState' effect to a 'State' effect, discarding
-- the notion of atomicity.
atomicStateToState :: Member (State s) r
                   => Sem (AtomicState s ': r) a
                   -> Sem r a
atomicStateToState :: Sem (AtomicState s : r) a -> Sem r a
atomicStateToState = (forall x (rInitial :: [(* -> *) -> * -> *]).
 AtomicState s (Sem rInitial) x -> Sem r x)
-> Sem (AtomicState s : r) a -> Sem r a
forall (e :: (* -> *) -> * -> *) (r :: [(* -> *) -> * -> *]) a.
FirstOrder e "interpret" =>
(forall x (rInitial :: [(* -> *) -> * -> *]).
 e (Sem rInitial) x -> Sem r x)
-> Sem (e : r) a -> Sem r a
interpret ((forall x (rInitial :: [(* -> *) -> * -> *]).
  AtomicState s (Sem rInitial) x -> Sem r x)
 -> Sem (AtomicState s : r) a -> Sem r a)
-> (forall x (rInitial :: [(* -> *) -> * -> *]).
    AtomicState s (Sem rInitial) x -> Sem r x)
-> Sem (AtomicState s : r) a
-> Sem r a
forall a b. (a -> b) -> a -> b
$ \case
  AtomicState f -> do
    (s
s', x
a) <- s -> (s, x)
f (s -> (s, x)) -> Sem r s -> Sem r (s, x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Sem r s
forall s (r :: [(* -> *) -> * -> *]).
MemberWithError (State s) r =>
Sem r s
get
    s -> Sem r ()
forall s (r :: [(* -> *) -> * -> *]).
MemberWithError (State s) r =>
s -> Sem r ()
put s
s'
    x -> Sem r x
forall (m :: * -> *) a. Monad m => a -> m a
return x
a
  AtomicState s (Sem rInitial) x
AtomicGet -> Sem r x
forall s (r :: [(* -> *) -> * -> *]).
MemberWithError (State s) r =>
Sem r s
get
{-# INLINE atomicStateToState #-}