State.hs

-- | Extend a monad with a modifiable environment
module Mini.Transformers.State (
  -- * Type
  StateT (
    StateT
  ),
  runStateT,

  -- * Operations
  get,
  modify,
  put,
) where

import Control.Applicative (
  Alternative,
  empty,
  (<|>),
 )
import Control.Monad (
  ap,
  liftM,
 )
import Control.Monad.IO.Class (
  MonadIO,
  liftIO,
 )
import Data.Bifunctor (
  first,
 )
import Mini.Random.Class (
  Random,
  random,
 )
import Mini.Transformers.Class (
  MonadTrans,
  lift,
 )
import Prelude (
  Applicative,
  Functor,
  Monad,
  MonadFail,
  const,
  fail,
  fmap,
  pure,
  ($),
  (.),
  (<$>),
  (<*>),
  (>>=),
 )

-- Type

-- | A transformer with state /s/, inner monad /m/, return /a/
newtype StateT s m a = StateT
  { runStateT :: s -> m (a, s)
  -- ^ Unwrap a transformer computation with an initial state
  }

instance (Monad m) => Functor (StateT s m) where
  fmap = liftM

instance (Monad m) => Applicative (StateT s m) where
  pure a = StateT $ \s -> pure (a, s)
  (<*>) = ap

instance (Monad m, Alternative m) => Alternative (StateT s m) where
  empty = StateT $ const empty
  m <|> n = StateT $ \s -> runStateT m s <|> runStateT n s

instance (Monad m) => Monad (StateT s m) where
  m >>= k = StateT $ \s -> runStateT m s >>= (\(a, s') -> runStateT (k a) s')

instance MonadTrans (StateT s) where
  lift m = StateT $ \s -> (\a -> (a, s)) <$> m

instance (MonadFail m) => MonadFail (StateT s m) where
  fail = StateT . const . fail

instance (MonadIO m) => MonadIO (StateT s m) where
  liftIO = lift . liftIO

instance (Monad m, Random a) => Random (StateT s m a) where
  random = first pure . random

-- Operations

-- | Fetch the current state
get :: (Monad m) => StateT s m s
get = StateT $ \s -> pure (s, s)

-- | Update the current state with an operation
modify :: (Monad m) => (s -> s) -> StateT s m ()
modify f = StateT $ \s -> pure ((), f s)

-- | Overwrite the current state with a value
put :: (Monad m) => s -> StateT s m ()
put s = StateT . const $ pure ((), s)