content: {-# 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 = send . 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' f = atomicState $ \s -> let (!s', a) = f s in (s', 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 = send AtomicGet
  {-# INLINE atomicGet #-}
  
  atomicGets :: Eff (AtomicState s) m => (s -> a) -> m a
  atomicGets = (<$> 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 f = atomicState $ \s -> (f 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' f = atomicState $ \s -> let !s' = f s in (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 = atomicState $ \_ -> (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 sInit main = do
    ref  <- embed $ newIORef sInit
    a    <- runAtomicStateIORef ref main
    sEnd <- embed $ readIORef ref
    return (sEnd, 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 sInit main = do
    ref  <- embed $ newIORef sInit
    a    <- runAtomicStateIORefSimple ref main
    sEnd <- embed $ readIORef ref
    return (sEnd, 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 ref = interpret $ \case
    AtomicState f -> embed (atomicModifyIORefP ref f)
    AtomicGet     -> embed (readIORef 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 ref = interpretSimple $ \case
    AtomicState f -> embed (atomicModifyIORefP ref f)
    AtomicGet     -> embed (readIORef 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 = 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 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 = interpretSimple $ \case
    AtomicState f -> embed $ atomically $ do
      (s, a) <- f <$> readTVar tvar
      writeTVar tvar s
      return a
    AtomicGet     -> embed (readTVarIO tvar)
  {-# INLINE runAtomicStateTVarSimple #-}
  
  data AtomicStateToStateH
  
  type AtomicStateToStateC s = InterpretC AtomicStateToStateH (AtomicState s)
  
  instance Eff (State s) m
        => Handler AtomicStateToStateH (AtomicState s) m where
    effHandler = \case
      AtomicState f -> state f
      AtomicGet     -> 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 = interpretViaHandler
  {-# INLINE atomicStateToState #-}
  