State Monad Transformer

(Not updated very recently)

StateT

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

instance Monad m => Monad (StateT s m) where
 -- return :: a -> StateT s m a
    return a = StateT $ \s -> pure (a, s)

 -- (>>=) :: StateT s m a -> (a -> StateT s m b) -> StateT s m b
    action >>= f = StateT $ \s ->
      do {- m -}
        (a, s1) <- runStateT action s0
        (b, s2) <- runStateT (f a) s1
        pure (b, s2)

Free instances:

instance (Monad m) => Functor (StateT s m) where {fmap f x = pure f <*> x}
instance (Monad m) => Applicative (StateT s m) where {pure = return; (<*>) = ap}

Plug in lift:

instance MonadTrans (StateT s) where
 -- lift :: Monad m => m a -> StateT s m a
    lift action = StateT $ \s -> flip (,) s <$> action

Add monoid structure:

instance (Monad m, Alternative m) => Alternative (StateT s m) where
 -- empty :: StateT s m a
    empty = StateT $ \s -> empty

 -- (<|>) :: StateT s m a -> StateT s m a -> StateT s m a
    x <|> y = StateT $ \s -> runStateT x s <|> runStateT y s

Helper functions:

get :: Monad m => StateT s m s
get = StateT $ \s -> pure (s, s)

put :: Monad m => s -> StateT s m ()
put s' = StateT $ \_ -> pure ((), s')

modify :: Monad m => (s -> s) -> StateT s m ()
modify f = StateT $ \s -> pure ((), f s)

evalStateT :: Monad m => StateT s m a -> s -> m a
evalStateT action init = fst <$> runStateT action init

execStateT :: Monad m => StateT s m a -> s -> m s
execStateT action init = snd <$> runStateT action init

Parser (Redux)

type Parser a = StateT String [] a

char :: Char -> Parser Char
char c = StateT $ \s ->
  case s of
    []     -> []
    (a:as) -> [(c,as) | a == c]

Only difference compared to our Parser library before is the use of the StateT constructor instead of Parser. And now we get extra functionality, such as get, push, and modify.

> runStateT (char 'a') "aabbcc"
[('a',"abbcc")]
> runStateT (char 'a' >> get) "aabbcc"
[("abbcc","abbcc")]

State (Redux)

type State s a = StateT s Identity a

Recall the trivial Identity wrapper monad.

newtype Identity a = Identity { runIdentity :: a }

instance Monad Identity where
    return  = Identity
    x >>= f = f $ runIdentity x

instance Functor Identity where {fmap f x = pure f <*> x}
instance Applicative Identity where {pure = return; (<*>) = ap}

Helper functions:

evalState :: State s a -> s -> a
evalState action init = fst $ runIdentity $ runStateT action init

execState :: State s a -> s -> s
execState action init = snd $ runIdentity $ runStateT action init

state :: (s -> (a, s)) -> State s a
state f = StateT $ \s -> pure $ f s

Back to Stack example. Same as before, since state and (>>=) take care of the pesky Identity wrappers:

type Stack = [Int]

push :: Int -> State Stack ()
pop :: State Stack Int

push i = state $ \stk -> ((), i:stk)

pop = state $ \stk ->
  case stk of
    i:is -> (i, is)
    []   -> (0, [])

Or:

push i = do
  is <- get
  put (i:is)

pop = do
  stk <- get
  case stk of
    []   -> pure 0
    i:is -> do
      put is
      pure i

NState

“Non-deterministic” stateful computations. Lists of possible results, try them all. Generalization of Parser with different types of state objects s.

type NState s a = StateT s [] a

Example:

type Turn = Either () ()
left      = Left ()
right     = Right ()

type Path = [Turn]

First an action that non-deterministically adds a step (i.e. adds both steps) to an exist Path:

step :: NState Path ()
step =
  StateT $ \path ->
    [ ((), dir:path) | dir <- [left, right] ]

Or:

step :: NState Path ()
step = modify (left:) <|> modify (right:)

Now take multiple steps:

walk :: NState Path ()
walk =                         -- walk = do
  step *> step *> step         --    step
                               --    step
                               --    step

runWalk :: [Path]
runWalk =
  execStateT walk []

> runWalk
... 8 possible paths ...

NState with Randomness

NState computations represent multiple potential results. What if we want to really non-deterministically (i.e. randomly) choose just one?

NState1

type NState1 s a = State (s, StdGen) a
             --  = StateT (s, StdGen) Identity a
             -- ~= (s, StdGen) -> (a, (s, StdGen))

First define a function that converts an “all possible results” computation into “one random result”:

chooseRandom1 :: NState s a -> NState1 s a
chooseRandom1 action = StateT $ \(s,g) ->
  let
    results = runStateT action s
    (i, g') = randomR (0, length results - 1) g
    (a, s') = results !! i
  in
  Identity $ (a, (s', g'))

And now perform a single random walk:

randomWalk1 :: NState1 Path ()
randomWalk1 =
  chooseRandom1 step
    *> chooseRandom1 step
    *> chooseRandom1 step

runRandomWalk1 :: IO Path
runRandomWalk1 = do
  g <- newStdGen
  pure $ fst $ execState randomWalk1 ([], g)

Or:

runRandomWalk1 =
  newStdGen >>= (,) [] >>> execState randomWalk1 >>> fst >>> pure

> runRandomWalk1
> runRandomWalk1
> runRandomWalk1

NState01

A downside of our approach is that chooseRandom1 will crash if results == []. A version that handles errors more gracefully:

type NState01 s a = StateT (s, StdGen) Maybe a
              -- ~= (s, StdGen) -> Maybe (a, (s, StdGen))

chooseRandom01 :: NState s a -> NState01 s a
chooseRandom01 action = StateT $ \(s,g) -> do
  let results = runStateT action s
  guard $ length results > 0
  let (i, g') = randomR (0, length results - 1) g
  let (a, s') = results !! i
  Just (a, (s', g'))

randomWalk01 :: NState01 Path ()
randomWalk01 = do
  chooseRandom01 step
    *> chooseRandom01 step
    *> chooseRandom01 step

runRandomWalk01 :: IO (Maybe Path)
runRandomWalk01 =
  newStdGen >>= (,) [] >>> execStateT randomWalk01 >>> fmap fst >>> pure

> runRandomWalk01
> runRandomWalk01
> runRandomWalk01

Library Version

Control.Monad.State.Lazy defines StateT and State.

Source Files

• StateT.hs