{-# LANGUAGE Trustworthy #-}
-- |
-- Copyright: (c) 2021 Xy Ren
-- License: BSD3
-- Maintainer: xy.r@outlook.com
-- Stability: experimental
-- Portability: non-portable (GHC only)
module Cleff.State
  ( -- * Effect
    State (..)
    -- * Operations
  , get
  , put
  , state
  , gets
  , modify
    -- * Interpretations
  , runState
  , runStateLocal
  , runStateIORef
  , runStateMVar
  , runStateTVar
  , zoom
  ) where

import           Cleff
import           Cleff.Internal.Base
import           Control.Concurrent.MVar (MVar, modifyMVar, readMVar, swapMVar)
import           Control.Concurrent.STM  (TVar, atomically, readTVar, readTVarIO, writeTVar)
import           Control.Monad           (void)
import           Data.Atomics            (atomicModifyIORefCAS)
import           Data.IORef              (IORef, newIORef, readIORef, writeIORef)
import           Data.ThreadVar          (getThreadVar, newThreadVar)
import           Data.Tuple              (swap)
import           Lens.Micro              (Lens', (&), (.~), (^.))

-- * Effect

-- | An effect capable of providing a mutable state @s@ that can be read and written. This roughly corresponds to the
-- @MonadState@ typeclass and @StateT@ monad transformer in the @mtl@ library.
data State s :: Effect where
  Get :: State s m s
  Put :: s -> State s m ()
  State :: (s -> (a, s)) -> State s m a

-- * Operations

makeEffect_ ''State

-- | Read the current state.
get :: State s :> es => Eff es s

-- | Update the state with a new value.
put :: State s :> es => s -> Eff es ()

-- | Modify the state /and/ produce a value from the state via a function.
state :: State s :> es
  => (s -> (a, s)) -- ^ The function that takes the state and returns a result value together with a modified state
  -> Eff es a

-- | Apply a function to the result of 'get'.
gets :: State s :> es => (s -> t) -> Eff es t
gets = (<$> get)

-- | Modify the value of the state via a function.
modify :: State s :> es => (s -> s) -> Eff es ()
modify f = state (((), ) . f)

-- * Interpretations

handleIORef :: IOE :> es => IORef s -> Handler (State s) es
handleIORef rs = \case
  Get     -> liftIO $ readIORef rs
  Put s'  -> liftIO $ writeIORef rs s'
  State f -> liftIO $ atomicModifyIORefCAS rs (swap . f)

-- | Run the 'State' effect.
--
-- === Caveats
--
-- The 'runState' interpreter is implemented with 'Data.IORef.IORef's and there is no way to do arbitrary
-- atomic transactions. The 'state' operation is atomic though and it is implemented with 'atomicModifyIORefCAS', which
-- can be faster than @atomicModifyIORef@ in contention. For any more complicated cases of atomicity, please build your
-- own effect that uses either @MVar@s or @TVar@s based on your need.
--
-- Unlike @mtl@, in @cleff@ the state /will not revert/ when an error is thrown.
--
-- 'runState' will stop taking care of state operations done on forked threads as soon as the main thread finishes its
-- computation. Any state operation done /before main thread finishes/ is still taken into account.
runState :: s -> Eff (State s : es) a -> Eff es (a, s)
runState s m = thisIsPureTrustMe do
  rs <- liftIO $ newIORef s
  x <- reinterpret (handleIORef rs) m
  s' <- liftIO $ readIORef rs
  pure (x, s')

-- | Run a 'State' effect where each thread has its thread-local state.
--
-- This means that each thread will have an individual state that has the same initial value. Threfore, state
-- operations on one thread will not change the state for any other thread.
--
-- The returned final state is that of the current thread.
--
-- === Caveats
--
-- Like 'runState', the 'state' operation in this handler is atomic. Like 'runState', and unlike @mtl@, any errors will
-- not revert the state changes.
--
-- Be warned that if you use a thread pool, then when a thread is reused, it may read the state left from the last
-- usage, therefore losing locality. If you use a thread pool, you will want to manually reset the state after each
-- task.
--
-- @since 0.3.3.0
runStateLocal :: s -> Eff (State s : es) a -> Eff es (a, s)
runStateLocal s m = thisIsPureTrustMe do
  rs <- liftIO $ newThreadVar s
  x <- reinterpret (\e -> liftIO (getThreadVar rs) >>= \r -> handleIORef r e) m
  s' <- liftIO $ readIORef =<< getThreadVar rs
  pure (x, s')

-- | Run the 'State' effect in terms of operations on a supplied 'IORef'. The 'state' operation is atomic.
--
-- @since 0.2.1.0
runStateIORef :: IOE :> es => IORef s -> Eff (State s : es) a -> Eff es a
runStateIORef rs = interpret $ handleIORef rs

-- | Run the 'State' effect in terms of operations on a supplied 'MVar'.
--
-- @since 0.2.1.0
runStateMVar :: IOE :> es => MVar s -> Eff (State s : es) a -> Eff es a
runStateMVar rs = interpret \case
  Get     -> liftIO $ readMVar rs
  Put s'  -> liftIO $ void $ swapMVar rs s'
  State f -> liftIO $ modifyMVar rs \s -> let (x, !s') = f s in pure (s', x)

-- | Run the 'State' effect in terms of operations on a supplied 'TVar'.
--
-- @since 0.2.1.0
runStateTVar :: IOE :> es => TVar s -> Eff (State s : es) a -> Eff es a
runStateTVar rs = interpret \case
  Get -> liftIO $ readTVarIO rs
  Put s' -> liftIO $ atomically $ writeTVar rs s'
  State f -> liftIO $ atomically do
    s <- readTVar rs
    let (x, !s') = f s
    writeTVar rs s'
    pure x

-- | Run a 'State' effect in terms of a larger 'State' via a 'Lens''.
zoom :: State t :> es => Lens' t s -> Eff (State s : es) ~> Eff es
zoom field = interpret \case
  Get     -> gets (^. field)
  Put s   -> modify (& field .~ s)
  State f -> state \t -> let (a, !s) = f (t ^. field) in (a, t & field .~ s)