{-# LANGUAGE BangPatterns, TupleSections #-}

module GStateT
    ( GStateT(..)
    , GPair(..)
    , Swap(..)
    , StateT
    , SwapStateT
    , runStateT
    , runSwapStateT
    , evalStateT
    , execStateT
    , mapStateT
    , get
    , put
    , modify
    , modify'
    ) where

import Control.Monad
import Control.Monad.Trans.Class (MonadTrans, lift)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Bifunctor
import Data.Tuple (swap)

class Bifunctor p => GPair p where
    pfst :: p a b -> a
    psnd :: p a b -> b
    both :: p a b -> (a, b)
    pair :: a -> b -> p a b

instance GPair (,) where
    pfst = fst
    psnd = snd
    both = id
    pair = (,)

-----------

newtype Swap a b = Swap { getSwap :: (b, a) } deriving (Eq, Show)

instance GPair Swap where
    pfst = snd . getSwap
    psnd = fst . getSwap
    both = swap . getSwap
    pair = (Swap .) . flip (,)

instance Bifunctor Swap where
    first f   = pair <$> f . pfst <*> psnd
    second f  = pair <$> pfst <*> f . psnd
    bimap f g = pair <$> f . pfst <*> g . psnd

instance Functor (Swap a) where
    fmap = second

instance (Ord a, Ord b) => Ord (Swap a b) where
    compare a b = compare (both a) (both b)

-----------

newtype GStateT p s m a = StateT { runGStateT :: s -> m (p a s) }

instance (GPair p, Monad m) => Functor (GStateT p s m) where
  fmap f m = StateT $ fmap (first f) . runGStateT m

instance (GPair p, Monad m) => Applicative (GStateT p s m) where
  pure x = StateT $ return . pair x
  f <*> m =
    StateT $ runGStateT f >=> fmap <$> first . pfst
                                  <*> runGStateT m . psnd

instance (GPair p, Monad m) => Monad (GStateT p s m) where
  return = pure
  m >>= f =
    StateT $ runGStateT m >=> runGStateT <$> f . pfst <*> psnd

instance GPair p => MonadTrans (GStateT p s) where
  lift m = StateT $ \ s -> flip pair s <$> m

instance (GPair p, MonadIO m) => MonadIO (GStateT p s m) where
  liftIO io = StateT $ \ s -> flip pair s <$> liftIO io

-------

type StateT = GStateT (,)
type SwapStateT = GStateT Swap

runStateT :: Monad m => StateT s m a -> s -> m (a, s)
runStateT = runGStateT

runSwapStateT :: Monad m => SwapStateT s m a -> s -> m (s, a)
runSwapStateT m = fmap getSwap . runGStateT m

evalStateT :: (GPair p, Functor m) => GStateT p s m a -> s -> m a
evalStateT m = fmap pfst . runGStateT m

execStateT :: (GPair p, Functor m) => GStateT p s m a -> s -> m s
execStateT m = fmap psnd . runGStateT m

-------

get :: (GPair p, Monad m) => GStateT p a m a
get = StateT $ return . (pair <$> id <*> id)

put :: (GPair p, Monad m) => s -> GStateT p s m ()
put = StateT . const . return . pair ()

modify :: (GPair p, Monad m) => (s -> s) -> GStateT p s m ()
modify f = StateT $ return . pair () . f

modify' :: (GPair p, Monad m) => s -> GStateT p s m ()
modify' f = get >>= \ s -> put $! s

mapStateT :: GPair p
          => (m (p a s) -> n (p b s))
          -> GStateT p s m a
          -> GStateT p s n b
mapStateT f m = StateT $ f . runGStateT m

withStateT :: (s -> s) -> GStateT p s m a -> GStateT p s m a
withStateT f m = StateT $ runGStateT m . f