module State where

import MonadIO
import MonadTrans

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

-- 1
instance (Functor m)
      => Functor (StateT s m) where
    fmap f (StateT m) = StateT $ \s -> fmap first $ m s
            where first = uncurry (\t1 t2 -> ((f t1), t2))
            
-- 2
-- Links:
--   http://stackoverflow.com/questions/18673525/is-it-possible-to-implement-applicative-m-applicative-statet-s-m
--   https://github.com/NICTA/course/issues/134
instance (Monad m)
      => Applicative (StateT s m) where
    pure x = StateT $ (\s -> pure (x, s))
    StateT g <*> StateT h = StateT $ \s -> keepFirst <$> g s <*> h s
        where keepFirst (f, s') (x, _) = (f x, s')
  

-- 3
instance (Monad m)
      => Monad (StateT s m) where
    return = pure

    (StateT sma) >>= f =
        StateT $ \s -> do
            a <- sma s
            runStateT (f $ fst a) s

instance MonadTrans (StateT s) where
    lift c = StateT $ \s -> c >>= (\x -> return (x, s))

instance (MonadIO m)
  => MonadIO (StateT s m) where
      liftIO = lift . liftIO