impl.hs

{-# LANGUAGE GADTs, FlexibleInstances, MultiParamTypeClasses #-}
import Control.Monad.State.Class
import Control.Monad.Trans.Class

data StateD s m a where
    Get :: StateD s m s
    Put :: s -> StateD s m ()
    Return :: a -> StateD s m a
    Bind :: StateD s m a -> (a -> StateD s m b) -> StateD s m b
    Lift :: m a -> StateD s m a

instance Monad (StateD s m) where
    return a = Return a
    ma >>= f = Bind ma f

instance MonadState s (StateD s m)  where
    get = Get
    put s = Put s

instance MonadTrans (StateD s) where
    lift = Lift

run :: Monad m => StateD s m a -> s -> m (s, a)
run Get s = return (s, s)
run (Put s') s = return (s', ())
run (Bind ma f) s = (run ma s) >>= (\(s',a) -> run (f a) s')
run (Return a) s = return (s, a)
run (Lift ma) s = ma >>= (\a -> return (s, a))

{-# RULES
"hacks" forall s. put s = Put s
"hacks2" get = Get
"put/put" forall s s'. Put s' >> Put s = Put s
"put/get" forall s. (Put s) >> Get = Put s >> return s
"get/put" Get >>= Put = return ()
  #-}

testAction :: StateD Int IO Int
testAction = do
  put 1
  put 2
  get

main = do
  (a, s) <- run testAction 1
  print (a,s)