Control.RMonad.Trans.State

{-# LANGUAGE CPP,
             MultiParamTypeClasses,
             NoImplicitPrelude,
             FlexibleContexts,
             FlexibleInstances,
             RebindableSyntax,
             TypeFamilies,
             UndecidableInstances #-}

module Control.RMonad.Trans.State where

import Control.RMonad.Prelude
import Control.RMonad

import Data.Suitable

-- anything suitable in the parent should be suitable here
newtype StateT s 𝔪 α = StateT { runStateT :: s -> 𝔪 (α, s) }

#define TYP_CONSTR (Suitable 𝔪 α, Suitable 𝔪 (α, s))
data instance Constraints (StateT s 𝔪) α =
    TYP_CONSTR => StateT_Constraints
instance TYP_CONSTR => Suitable (StateT s 𝔪) α where
    constraints = StateT_Constraints

instance RMonad 𝔪 => RMonad (StateT s 𝔪) where
    return x = withResConstraints $ \StateT_Constraints ->
        (StateT $ \s -> return (x, s))
    m >>= k = withConstraintsOf m $ \StateT_Constraints ->
        withResConstraints $ \StateT_Constraints ->
        StateT $ \s -> do
            ~(a, s') <- runStateT m s
            runStateT (k a) s'

-- |Similar to 'evalState'
evalStateT
  :: (Suitable m (b, s), Suitable m b, RMonad m) =>
     StateT s m b -> s -> m b
evalStateT m s = do
    ~(a, _) <- runStateT m s
    return a

-- |Similar to 'execState'
execStateT
  :: (Suitable m (α, b), Suitable m b, RMonad m) =>
     StateT b m α -> b -> m b
execStateT m s = do
    ~(_, s') <- runStateT m s
    return s'

-- |Similar to 'mapState'
mapStateT
  :: (𝔪1 (α1, s) -> 𝔪 (α, s)) -> StateT s 𝔪1 α1 -> StateT s 𝔪 α
mapStateT f m = StateT $ f . runStateT m

-- |Similar to 'withState'
withStateT :: (s -> s) -> StateT s 𝔪 α -> StateT s 𝔪 α
withStateT f m = StateT $ runStateT m . f

instance (RFunctor m) => RFunctor (StateT s m) where
    fmap f m = withConstraintsOf m $ \StateT_Constraints ->
        withResConstraints $ \StateT_Constraints ->
        StateT $ \ s ->
            fmap (\ ~(a, s') -> (f a, s')) $ runStateT m s

-- NOTE: type inference will not infer this header
get :: (RMonad 𝔪, Suitable 𝔪 (s, s)) => StateT s 𝔪 s
get = StateT $ \s -> return (s, s)

-- | More useful than put, returns previous value
set :: (RMonad 𝔪, Suitable 𝔪 (s, s)) => s -> StateT s 𝔪 s
set s = StateT $ \s_old -> return (s_old, s)

-- | Monadic state transformer.
--      Maps an old state to a new state inside a state monad.
--      The old state is returned
modify :: (RMonad 𝔪, Suitable 𝔪 s, Suitable 𝔪 (s, s)) =>
    (s -> s) -> StateT s 𝔪 s
modify f = get >>= set . f

-- | Gets specific component of the state, using a projection function
-- supplied.
gets
  :: (RMonad 𝔪,
      Suitable 𝔪 (a, a),
      Suitable 𝔪 (b, a),
      Suitable 𝔪 b,
      Suitable 𝔪 a) =>
     (a -> b) -> StateT a 𝔪 b
gets f = get >>= return . f