Stateful Functions

Many tasks involve maintaining and “updating” a set of data that summarizes the “state” (or “model”) of the running program.

Example 1: Stacks

Let’s define the following to represent and manipulate stacks, a data structure that allows adding and removing element only from the top end:

type Stack a = [a]

push :: a -> Stack a -> Stack a
pop  :: Stack a -> (a, Stack)

push a as  = a:as
pop (a:as) = (a, as)

We can test our Stack with a sequence of operations:

testStack_ :: Stack Int -> (Int, Stack Int)
testStack_ s0 =
  let
    s1 = push 0 s0
    s2 = push 1 s1
    s3 = push 2 s2
    s4 = push 3 s3
    (a,s5) = pop s4
    (b,s6) = pop s5
    s7 = push (a+b) s6
  in
    pop s7

> testStack_ []
(5,[1,0])

The stack is “threaded through” a series of computations. Let’s make this even more explicit. Compared to…

push :: a -> Stack a ->      Stack a

… let’s redefine push to also produce a (the) () value:

push :: a -> Stack a -> ((), Stack a)
push a as = ((), a:as)

Now our sequence of operations looks like the following:

testStack :: Stack Int -> (Int, Stack Int)
testStack s0 =
  let
    (_,s1) = push 0     s0
    (_,s2) = push 1     s1
    (_,s3) = push 2     s2
    (_,s4) = push 3     s3
    (a,s5) = pop        s4
    (b,s6) = pop        s5
    (_,s7) = push (a+b) s6
  in
    pop s7

What’s wrong with this? First of all, there’s a lot of syntactic “noise.” Worse yet, it’s easy to make a mistake; the intention is that each version of the stack s_i is referred to once and then “discarded” in favor of the new stack returned by the operation.

“Stateful Functions”

There is a common idiom in the code above: the threading of values of some particular “state” type (referred to here as Model) through a series of computations.

computation :: Model -> (T7, Model)
computation s0 =
  let
    (a1,s1) = f0 s0    -- f0 :: Model -> (T1, Model)
    (a2,s2) = f1 s1    -- f1 :: Model -> (T2, Model)    f1 may use a1
    (a3,s3) = f2 s2    -- f2 :: Model -> (T3, Model)    f2 may use a{1,2}
    (a4,s4) = f3 s3    -- f3 :: Model -> (T4, Model)    f3 may use a{1,2,3}
    (a5,s5) = f4 s4    -- f4 :: Model -> (T5, Model)    ...
    (a6,s6) = f5 s5    -- f5 :: Model -> (T6, Model)
    (a7,s7) = f6 s6    -- f6 :: Model -> (T7, Model)
  in
    (a7,s7)

Note: I am purposely avoiding the intuitive name State — instead choosing the somewhat less helpful name Model — for reasons that will become clear shortly.

We refer to this object as “the state” even though, if you’re familiar with other languages with “mutable” or “stateful” features, there’s nothing like that here. Just ordinary pure functions, with a pattern of use that feels like we’re manipulating state.

Notice that each of the functions f_i can refer to the results of previous function calls. Where do these arguments, the ones besides the Model argument, appear? In general, a (curried) function of the form

(T_i, Model) -> (T_{i+1}, Model)

can be rewritten as the (uncurried) function

T_i -> Model -> (T_{i+1}, Model)

So, when identifying the essential recurring pattern above, it works out best to take the Model argument separately, and last.

type StateFunc s a =
    s -> (a,s)             -- name for function type idiom

data State s a =
    State (s -> (a,s))     -- new datatype to define instances
 -- State (StateFunc s a)

newtype State s a =
    State (s -> (a,s))     -- newtype b/c one, one-arg constructor
 -- State (StateFunc s a)

newtype State s a =
    State { runState :: s -> (a,s) }     -- field for unboxing
 -- State { runState :: StateFunc s a }

StateFunc s a means a computation that, starting with an input state of type s, produces a value of type a and an updated state of type s.

When you read State s a, think “stateful function of type s -> (a,s)” (but boxed up in a newtype). Or think “stateful computation that produces an a” keeping in mind that there is input and output state “in the background.”

Note: The name “State” for this type and its data constructor follow choices in the library. The type argument s is the type of the “model” or the “state.”

Where Are We Going With This?

To preview the benefits that this abstraction will provide, we are going to define

pop'  :: State (Stack a) a
push' :: a -> State (Stack a) ()

and, then, because State s will be a Monad, the previous sequence of stack operations becomes:

testStack' = do
  push' 0
  push' 1
  push' 2
  push' 3
  a <- pop'
  b <- pop'
  push' (a+b)
  pop'

Sequencing Stateful Functions

We can think of sequencing States as function composition, with the appropriate plumbing to thread the state objects through the component functions.

     s0   ------  s1   ------  s2   ------  s3
    ----> | f0 | ----> | f1 | ----> | f2 | ---->
          ------       ------       ------
              \_________/  \_________/  \------>
                a1           a2           a3

   f0                               :: s -> (a1, s)
   f1 {- f1 may use a1     -}       :: s -> (a2, s)
   f2 {- f2 may use a{1,2} -}       :: s -> (a3, s)
  
   f0 >>= \a1 -> f1 >>= \a2 -> f2   :: s -> (a3, s)

So, let’s define how State s forms a monadic, applicative functor.

fmap  :: (a -> b) -> State s a -> State s b
(<*>) :: State s (a -> b) -> State s a -> State s b
(>>=) :: State s a -> (a -> State s b) -> State s b
pure  :: a -> State s a

Let’s define the Monad instance, and then derive free instances for Functor and Applicative. (Alternatively, to develop the intuition for how stateful functions work more slowly, define the instances “in order”.)

instance Monad (State s) where
  return :: a -> State s a
  return a = State $ \s -> (a, s)

  (>>=) :: State s a -> (a -> State s b) -> State s b
  sa >>= f = State $ \s0 ->
    let
      (a, s1) = runState sa s0
      (b, s2) = runState (f a) s1
    in
      (b, s2)

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

Example: Stacks

pop'  :: State (Stack a) a
push' :: a -> State (Stack a) ()

pop'    = State $ \(a:as) -> (a,as)
push' a = State $ \as -> ((), a:as)

Now, let’s go back to our long sequence of stack operations.

testStack' :: State (Stack Int) Int
testStack' = do
  push' 0
  push' 1
  push' 2
  push' 3
  a <- pop'
  b <- pop'
  push' (a+b)
  pop'

> runState testStack' []
(5,[1,0])

Whoa, cool!

Expressions like push' 0 look like they’re missing something. Where did the stack go? It may help to note an identity regarding our previous version of push:

push' 0  ===  (\s0 -> push' 0 s0)

Helper Functions

get :: State s s                        -- get state out
get = State $ \s -> (s, s)

put :: s -> State s ()                  -- set "current" state
put s' = State $ \s -> ((), s')

modify :: (s -> s) -> State s ()        -- modify the state
modify f = State $ \s -> ((), f s)

evalState :: State s a -> s -> a        -- run and return final value
evalState sa s = fst $ runState sa s

execState :: State s a -> s -> s        -- run and return final state
execState sa s = snd $ runState sa s

Example: Stacks

Now, if we wanted to, we can redefine pop' and push' using the Monadic interface and the helpers for “reading” and “writing” the state (Stack) in the background.

push' a = do                -- push' a =
  as <- get                 --   get >>= \as ->
  put (a:as)                --   put (a:as)

pop' = do                   -- pop' =
  (a:as) <- get             --   get >>= \(a:as) ->
  put as                    --     put as >>
  pure a                    --     pure a

Example 2: Random Numbers

Let’s work through a second example.

Pure Pseudo-Random Number Generators

You may want to install random globally: cabal install --lib random
Otherwise, you may need a local .cabal file and: :set -package random

> import System.Random
> import System.Random.Stateful

> :t uniform
uniform :: (RandomGen g, Uniform a) => g -> (a, g)

Uniform: “The class of types for which a uniformly distributed value can be drawn from all possible values of the type.”
RandomGen: “An interface to pure pseudo-random number generators. StdGen is the standard RandomGen instance provided by this library.”

Okay, so how do we get a StdGen?

mkStdGen: “Constructs a StdGen deterministically.”

> :t mkStdGen
> let g = mkStdGen 17

> :t random g
> :t fst $ random g
> fst $ uniform g            -- by default, GHCi thinks we want ()

> fst $ uniform @StdGen @Int g
> fst $ uniform g :: Int
> fst $ uniform g :: Bool
> fst $ uniform g :: Float

Can we get multiple uniform numbers?

> fst $ uniform g :: Int
> fst $ uniform g :: Int
> fst $ uniform g :: Int

This is pseudo-randomness: a “sequence of numbers is one that appears to be statistically random, despite having been produced by a completely deterministic and repeatable process.”

Need to thread the StdGens through…

generateThree_ :: Uniform a => StdGen -> ((a,a,a), StdGen)
generateThree_ g0 =
  let
    (a1,g1) = uniform g0
    (a2,g2) = uniform g1
    (a3,g3) = uniform g2
  in
    ((a1, a2, a3), g3)

Look familiar?!? We’ll clean this up shortly.

But first, we don’t want to explicitly pick a StdGen on each run…

> fst $ generateThree_ @Int $ mkStdGen 17
(8546247003154199974,5724321524807575004,-2241565262936054436)

… we need to start the deterministic process with a non-deterministic choice.

Random Initial Seed

Want Haskell to (randomly) pick one for us.

> :t initStdGen @IO
initStdGen @IO :: IO StdGen

> g <- initStdGen
> g

initStdGn: “Initialize StdGen using system entropy (i.e. /dev/urandom) when it is available, while falling back on using system time as the seed.

Notice the IO scarlet letter; there is communication with the world. Try it out with do-notation…

> do { g <- initStdGen; pure $ fst $ generateThree_ @Int g }

… and then clean up:

> fst <$> generateThree_ @Int <$> initStdGen @IO
> fst <$> generateThree_ @Int <$> initStdGen @IO
> fst <$> generateThree_ @Int <$> initStdGen @IO

The `RandState` Monad

The type State StdGen a describes (wrapped) functions of type StdGen -> (a, StdGen).

type RandState a = State StdGen a
type RandState   = State StdGen

Start by writing a computation that produces a single value.

generateOne :: Uniform a => State StdGen a
generateOne :: Uniform a => RandState a

generateOne = State uniform

Or, can rewrite with do-notation. (Arguably more confusing, because uniform takes StdGen arg.)

generateOne = do
  g0 <- get                   -- get >>= \g0 ->
  let (a,g1) = uniform g0     -- let (a,g1) = uniform g0 in
  put g1                      -- put g1 >>
  pure a                      --   pure a

Now we can easily sequence calls together.

generateThree :: Uniform a => State StdGen (a, a, a)
generateThree :: Uniform a => RandState (a, a, a)

generateThree =
  generateOne >>= \a1 ->
  generateOne >>= \a2 ->
  generateOne >>= \a3 ->
    pure (a1, a2, a3)

Can rewrite with do-notation.

generateThree = do
  a1 <- generateOne
  a2 <- generateOne
  a3 <- generateOne
  pure (a1, a2, a3)

Better yet, in applicative style:

generateThree =
  liftA3 (,,) generateOne generateOne generateOne

> (evalState $ generateThree @Int) <$> initStdGen
> (evalState $ generateThree @Int) <$> initStdGen
> (evalState $ generateThree @Int) <$> initStdGen

Source Files

Misc

Library

The way that the State type is implemented in the Control.Monad.State library is more involved than our definition here. We’ll talk a little about the library approach later.

Random History

There are some recent changes to the System.Random libraries. If you come across examples written using the older interface (which has been kept around for backward compatibility), you can port things to the new interface as follows. See System.Random and System.Random.Stateful for more details.

Old	New
`Random`	`Uniform`
`random`	`uniform`
`randomR`	`uniformR`
`newStdGen` / `getStdGen`	`initStdGen @IO`
`randomRIO range`	`fst <$> uniformR range <$> initStdGen`

Applicative Effects

A not-free instance for Applicative might look like:

sf <*> sa = State $ \s0 ->
  let
    (f, s1) = runState sf s0
    (a, s2) = runState sa s1
  in
    (f a, s2)

Why not run the stateful functions in the opposite order? By convention, Applicative instances sequence effects from left-to-right, as described here:

Incidentally, the fact (>>=) performing left-to-right sequencing is the main reason for the convention that leads applicative operators to do the same. liftM2 and ap are implemented using (>>=), and so they also sequence effects from left to right. That means Applicative instances must follow suit if they are to be coherent with Monad ones, and at that point it becomes sensible to extend the convention to all applicative functors (even those without Monad instances) to minimise a source of confusion.

Example 1: Stacks

“Stateful Functions”

Where Are We Going With This?

Sequencing Stateful Functions

Example: Stacks

Helper Functions

Example: Stacks

Example 2: Random Numbers

Pure Pseudo-Random Number Generators

Random Initial Seed

The RandState Monad

Source Files

Misc

Library

Random History

Applicative Effects

The `RandState` Monad