{-# LANGUAGE TupleSections #-}
module Foundation.Monad.State
    ( -- * MonadState
      MonadState(..)
    , get
    , put

    , -- * StateT
      StateT
    , runStateT
    ) where

import Foundation.Internal.Base (($), (.), const)
import Foundation.Monad.Base

class Monad m => MonadState m where
    type State m
    withState :: (State m -> (a, State m)) -> m a

get :: MonadState m => m (State m)
get = withState $ \s -> (s, s)

put :: MonadState m => State m -> m ()
put s = withState $ const ((), s)

-- | State Transformer
newtype StateT s m a = StateT { runStateT :: s -> m (a, s) }

instance Functor m => Functor (StateT s m) where
    fmap f m = StateT $ \s1 -> (\(a, s2) -> (f a, s2)) `fmap` runStateT m s1
    {-# INLINE fmap #-}

instance (Applicative m, Monad m) => Applicative (StateT s m) where
    pure a     = StateT $ \s -> (,s) `fmap` pure a
    {-# INLINE pure #-}
    fab <*> fa = StateT $ \s1 -> do
        (ab,s2) <- runStateT fab s1
        (a, s3) <- runStateT fa s2
        return (ab a, s3)
    {-# INLINE (<*>) #-}

instance (Functor m, Monad m) => Monad (StateT s m) where
    return a = StateT $ \s -> (,s) `fmap` return a
    {-# INLINE return #-}
    ma >>= mab = StateT $ \s1 -> runStateT ma s1 >>= \(a, s2) -> runStateT (mab a) s2
    {-# INLINE (>>=) #-}

instance MonadTrans (StateT s) where
    lift f = StateT $ \s -> f >>= return . (,s)
    {-# INLINE lift #-}

instance (Functor m, MonadIO m) => MonadIO (StateT s m) where
    liftIO f = lift (liftIO f)
    {-# INLINE liftIO #-}

instance (Functor m, MonadFailure m) => MonadFailure (StateT s m) where
    type Failure (StateT s m) = Failure m
    mFail e = StateT $ \s -> ((,s) `fmap` mFail e)

instance (Functor m, MonadThrow m) => MonadThrow (StateT s m) where
    throw e = StateT $ \_ -> throw e

instance (Functor m, MonadCatch m) => MonadCatch (StateT s m) where
    catch (StateT m) c = StateT $ \s1 -> m s1 `catch` (\e -> runStateT (c e) s1)

instance (Functor m, Monad m) => MonadState (StateT s m) where
    type State (StateT s m) = s
    withState f = StateT $ return . f