content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE TemplateHaskell #-}
  
  -- | Description: The 'AtomicState' effect
  module Polysemy.AtomicState
    ( -- * Effect
      AtomicState (..)
  
      -- * Actions
    , atomicState
    , atomicState'
    , atomicGet
    , atomicGets
    , atomicPut
    , atomicModify
    , atomicModify'
  
      -- * Interpretations
    , runAtomicStateIORef
    , runAtomicStateTVar
    , atomicStateToIO
    , atomicStateToState
    , runAtomicStateViaState
    , evalAtomicStateViaState
    , execAtomicStateViaState
    ) 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
    -- | Run a state action.
    AtomicState :: (s -> (s, a)) -> AtomicState s m a
    -- | Get the state.
    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 = (<$> 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' 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 <- atomicState $ \s ->
      case f s of
        v@(!_, _) -> v
    return a
  {-# INLINE atomicState' #-}
  
  -----------------------------------------------------------------------------
  -- | Replace the state with the given value.
  atomicPut :: Member (AtomicState s) r
            => s
            -> Sem r ()
  atomicPut s = do
    !_ <- atomicState $ \_ -> (s, ()) -- strict put with atomicModifyIORef
    return ()
  {-# INLINE atomicPut #-}
  
  -----------------------------------------------------------------------------
  -- | Modify the state lazily.
  atomicModify :: Member (AtomicState s) r
               => (s -> s)
               -> Sem r ()
  atomicModify f = atomicState $ \s -> (f 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' f = do
    !_ <- atomicState $ \s -> let !s' = f s in (s', ())
    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 ref = interpret $ \case
    AtomicState f -> embed $ atomicModifyIORef ref f
    AtomicGet     -> embed $ readIORef 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 = interpret $ \case
    AtomicState f -> embed $ atomically $ do
      (s', a) <- f <$> readTVar tvar
      writeTVar tvar s'
      return a
    AtomicGet -> embed $ readTVarIO 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 = do
    ref <- embed $ newIORef s
    res <- runAtomicStateIORef ref sem
    end <- embed $ readIORef ref
    return (end, 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 = interpret $ \case
    AtomicState f -> do
      (s', a) <- f <$> get
      put s'
      return a
    AtomicGet -> get
  {-# INLINE atomicStateToState #-}
  
  ------------------------------------------------------------------------------
  -- | Run an 'AtomicState' with local state semantics, discarding
  -- the notion of atomicity, by transforming it into 'State' and running it
  -- with the provided initial state.
  --
  --
  -- @since v1.7.0.0
  runAtomicStateViaState :: s
                         -> Sem (AtomicState s ': r) a
                         -> Sem r (s, a)
  runAtomicStateViaState s =
    runState s . atomicStateToState . raiseUnder
  {-# INLINE runAtomicStateViaState #-}
  
  ------------------------------------------------------------------------------
  -- | Evaluate an 'AtomicState' with local state semantics, discarding
  -- the notion of atomicity, by transforming it into 'State' and running it
  -- with the provided initial state.
  --
  -- @since v1.7.0.0
  evalAtomicStateViaState :: s
                          -> Sem (AtomicState s ': r) a
                          -> Sem r a
  evalAtomicStateViaState s =
    evalState s . atomicStateToState . raiseUnder
  {-# INLINE evalAtomicStateViaState #-}
  
  ------------------------------------------------------------------------------
  -- | Execute an 'AtomicState' with local state semantics, discarding
  -- the notion of atomicity, by transforming it into 'State' and running it
  -- with the provided initial state.
  --
  -- @since v1.7.0.0
  execAtomicStateViaState :: s
                          -> Sem (AtomicState s ': r) a
                          -> Sem r s
  execAtomicStateViaState s =
    execState s . atomicStateToState . raiseUnder
  {-# INLINE execAtomicStateViaState #-}
  