{- |
Strict & properly inline StateT.
-}

module Utils.StateT (
    module Utils.StateT
  , module Control.Monad.State.Class
  , module Control.Monad.Trans.Class
  ) where

import Control.Monad.Error.Class
import Control.Monad.IO.Class
import Control.Monad.Reader.Class
import Control.Monad.State.Class
import Control.Monad.Trans.Class
import Control.Monad.Writer.Class
import GHC.Exts (oneShot)

import Utils.Pair

newtype StateT s m a = StateT# {runStateT :: s -> m (Pair a s)}

pattern StateT :: (s -> m (Pair a s)) -> StateT s m a
pattern StateT f <- StateT# f where
  StateT f = StateT# (oneShot f)
{-# inline StateT #-}
{-# complete StateT #-}

instance Functor m => Functor (StateT s m) where
  {-# inline fmap #-}
  fmap = \f (StateT g) -> StateT (\s -> fmap (\(Pair a s) -> Pair (f a) s) (g s))
  {-# inline (<$) #-}
  (<$) = \a m -> (\_ -> a) <$> m

-- We require Monad m in order to force strictness in the Applicative sequencing.
instance Monad m => Applicative (StateT s m) where
  {-# inline pure #-}
  pure = \a -> StateT (\s -> pure (Pair a s))
  {-# inline (<*>) #-}
  (<*>) = \(StateT mf) (StateT ma) -> StateT (\s -> do
    Pair f s <- mf s
    Pair a s <- ma s
    let b = f a
    pure (Pair b s))
  {-# inline (*>) #-}
  (*>) = \(StateT ma) (StateT mb) -> StateT (\s -> do
     Pair _ s <- ma s
     mb s)
  {-# inline (<*) #-}
  (<*) = \(StateT ma) (StateT mb) -> StateT (\s -> do
    Pair a s <- ma s
    Pair _ s <- mb s
    pure (Pair a s))

instance Monad m => Monad (StateT s m) where
  {-# inline return #-}
  return = pure
  {-# inline (>>=) #-}
  (>>=) = \(StateT ma) f -> StateT (\s -> do
    Pair a s <- ma s
    runStateT (f a) s)
  {-# inline (>>) #-}
  (>>) = (*>)

instance MonadTrans (StateT s) where
  {-# inline lift #-}
  lift ma = StateT (\s -> do a <- ma; pure (Pair a s))

instance Monad m => MonadState s (StateT s m) where
  {-# inline state #-}
  state = \f -> StateT (\s -> case f s of (!a, !s) -> pure (Pair a s))
  {-# inline get #-}
  get = StateT (\s -> pure (Pair s s))
  {-# inline put #-}
  put = \s -> StateT (\_ -> pure (Pair () s))


instance MonadFail m => MonadFail (StateT s m) where
  {-# inline fail #-}
  fail = \str -> StateT \_ -> fail str

instance MonadIO m => MonadIO (StateT s m) where
  {-# inline liftIO #-}
  liftIO = \ma -> lift (liftIO ma)

instance MonadReader r m => MonadReader r (StateT s m) where
  {-# inline ask #-}
  ask = lift ask
  {-# inline local #-}
  local = \f (StateT ma) -> StateT (\s -> local f (ma s))

instance MonadWriter w m => MonadWriter w (StateT s m) where
  {-# inline writer #-}
  writer = lift . writer
  {-# inline tell #-}
  tell   = lift . tell
  {-# inline listen #-}
  listen = \m -> StateT (\s -> do
    ((a `Pair` s'),w) <- listen (runStateT m s)
    return $! ((a, w) `Pair` s'))
  {-# inline pass #-}
  pass = \m -> StateT (\s -> pass $ do
    ((a, f) `Pair` s') <- runStateT m s
    return $! ((a `Pair` s'), f))

instance MonadError e m => MonadError e (StateT s m) where
  {-# inline throwError #-}
  throwError = lift . throwError
  {-# inline catchError #-}
  catchError = \m h -> StateT (\s -> runStateT m s `catchError` \e -> runStateT (h e) s)

{-# inline execStateT #-}
execStateT :: Monad m => StateT s m a -> s -> m s
execStateT = \(StateT f) s -> do _ `Pair` s <- f s; pure s

{-# inline evalStateT #-}
evalStateT :: Monad m => StateT s m a -> s -> m a
evalStateT = \(StateT f) s -> do a `Pair` _ <- f s; pure a