module Control.Monad.Stack (
      Stack
    , runStack
    , StackT
    , runStackT
    , peek
    , pop
    , push
    , peeks
    , testTop
    ) where

import Data.Maybe
import Control.Applicative
import Control.Monad
import Control.Monad.Identity
import Control.Monad.State.Strict
import Control.Monad.Trans


type Stack s = StackT s Identity

newtype StackT s m a = StackT { unStack :: StateT [s] m a }

runStack :: Stack s a -> a
runStack = runIdentity . runStackT

runStackT :: (Monad m) => StackT s m a -> m a
runStackT = flip evalStateT [] . unStack


peek :: (Monad m) => StackT s m (Maybe s)
peek = StackT $ do
    stack <- get
    return $ case stack of
        [] -> Nothing
        (x:_) -> Just x
pop :: (Monad m) => StackT s m (Maybe s)
pop = StackT $ do
    stack <- get
    case stack of
        [] -> return Nothing
        (x:xs) -> put xs >> return (Just x)
push :: (Monad m) => s -> StackT s m ()
push x = StackT $ modify (x:) >> return ()


peeks :: (Monad m) => (s -> a) -> StackT s m (Maybe a)
peeks f = (fmap . fmap) f peek

testTop :: (Monad m) => (s -> Bool) -> StackT s m Bool
testTop = (fromMaybe False <$>) . peeks



instance (Monad m) => Functor (StackT s m) where
    fmap = liftM
instance (Monad m) => Applicative (StackT s m) where
    pure = return
    (<*>) = ap

instance (Monad m) => Monad (StackT s m) where
    return = StackT . return
    x >>= k = StackT $ unStack x >>= unStack . k

instance MonadTrans (StackT s) where
    lift = StackT . lift

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