{-# LANGUAGE TupleSections #-}

-- | Extension of a monad with a modifiable environment
module Mini.Transformers.StateT (
  -- * Type
  StateT,

  -- * Reading
  get,

  -- * Modifying
  modify,

  -- * Writing
  put,

  -- * Runners
  runStateT,
) where

import Control.Applicative (
  Alternative (
    empty,
    (<|>)
  ),
 )
import Control.Monad (
  ap,
  liftM,
  (>=>),
 )
import Data.Functor (
  (<&>),
 )
import Mini.Transformers.Class (
  MonadTrans (
    lift
  ),
 )

{-
 - Type
 -}

-- | A monad with modifiable type /s/, inner monad /m/, and return type /a/
newtype StateT s m a = StateT
  { runStateT :: s -> m (a, s)
  -- ^ Unwrap a 'StateT' given a starting 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 $ runStateT m >=> (\(a, s) -> runStateT (k a) s)

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

{-
 - Reading
 -}

{- | Fetch the current value of the state

> foo = do
>   s <- get
>   bar s
-}
get :: (Monad m) => StateT s m s
get = StateT $ \s -> pure (s, s)

{-
 - Modifying
 -}

{- | Update the current value of the state with the given operation

> foo = modify $ \s -> bar s
-}
modify :: (Monad m) => (s -> s) -> StateT s m ()
modify f = StateT $ pure . ((),) . f

{-
 - Writing
 -}

-- | Set the state to the given value
put :: (Monad m) => s -> StateT s m ()
put = StateT . const . pure . ((),)