State-solution.hs

{-# LANGUAGE RankNTypes #-}

newtype State s a = State { runState :: forall r. s -> (a -> s -> r) -> r }

stateReturn :: a -> State s a
stateReturn a = State $ \s k -> k a s

stateMap :: (a -> b) -> State s a -> State s b
stateMap f (State sa) = State $ \s kb -> sa s (kb . f)

stateBind :: (a -> State s b) -> State s a -> State s b
stateBind f (State sa) = State $ \s kb -> sa s (\a s' -> runState (f a) s' kb)

stateGet :: State s s
stateGet = State $ \s k -> k s s

statePut :: s -> State s ()
statePut s = State $ \_ k -> k () s

stateModify  :: (s -> s) -> State s ()
stateModify f = stateBind (statePut . f) stateGet

execState :: s -> State s a -> s
execState s st = runState st s (flip const)

evalState :: s -> State s a -> a
evalState s st = runState st s const