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

    , -- * StateT
      StateT
    , runStateT
    ) where

import Basement.Compat.Bifunctor (first)
import Basement.Compat.Base (($), (.), const)
import Foundation.Monad.Base
import Control.Monad ((>=>))

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 -> (first f) `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 $ runStateT ma >=> (\(a, s2) -> runStateT (mab a) s2)
    {-# INLINE (>>=) #-}

instance (Functor m, MonadFix m) => MonadFix (StateT s m) where
    mfix f = StateT $ \s -> mfix $ \ ~(a, _) -> runStateT (f a) s
    {-# INLINE mfix #-}

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