home/blog/supplements/monads/state-example

State usage example

State can be a bit confusing at first. Here’s a small usage example. As a reminder, here are the relevant definitions from the main post.

data State s a = State (s -> (s, a))

instance Functor (State s) where
  fmap f (State proc) = State proc'
    where proc' = applySnd f . proc

applySnd :: (a -> b) -> (s, a) -> (s, b)
applySnd f (state, value) = (state, f value)

unwrapState :: State s a -> (s -> (s, a))
unwrapState (State proc) = proc

instance Monad (State s) where
  return x = State (\s -> (s, x))

  bind f (State proc) = State proc'
    where f' = unwrapState . f
      proc' s0 = (s2, val'')
        where (s1, val1) = proc s0
              (s2, val2) = f' val1 s1

The example we’ll use is random number generation. In many other languages, sampling a random value is a fairly trivial thing: Usually, it’s either a function that takes no arguments:

def random() -> int:
  # generate a random integer

or it’s implemented as a method on a type that mutates some internal data:

struct Generator { /* hidden data */ }

impl Generator {
    fn sample(&mut self) -> i32 { /* generate a random integer */ }
}

Due to how generation works this is actually a stateful computation, and any implementation of RNG using pseudorandom sequences must account for it. In Haskell, direct mutation of values is, of course, not allowed, which implies that RNG needs to be a bit more complicated. Specifically, it looks like a type to store a location in a pseudorandom sequence and a function to compute a value and the next location in the sequence. Critically, however, must be included in the return value of the sampling function per the rules of Haskell.

data Generator = -- represents a location in a pseudorandom sequence

-- use a seed to initialize
makeGenerator :: Int -> Generator
makeGenerator seed = -- some function of the seed...

-- generate a random integer and return the new point in the sequence
sample :: Generator -> (Int, Generator)
sample gen = (num, gen')
  where (num, gen') = -- some simultaneous process for both the number
                      -- generation and the evolution of the sequence

This makes the bare sampling function annoying to use on its own, since doing so would require always explicitly returning the new representation of the sequence from every function that uses it. Instead, we’ll wrap it in a State:

type RandomInt = State Generator Int

At this point, it helps to introduce two functions to easily interact with the state value within a monadic context. These are known as get and put:

-- `get` replaces the output value with the state value
get :: State s s
get = State copy

-- return two copies of a value
copy :: s -> (s, s)
copy x = (x, x)

-- `put` replaces both the output and state value with a new state value
put :: s -> State s s
pub x = State (replace x)

-- disregard the second argument, taking only the first
replace :: s -> s -> (s, s)
replace x _ = (x, x)

It’s worth noting that these functions’ signatures may look a little weird at first–we think of them as performing manipulations on a State, so why are they typed like ordinary values? The answer lies in the fact that State is not really a value, it’s actually more of a transformation between values; the actual manipulation of the state is only realized when we compose these transformations together.

Using get and put, we can write a more convenient function to sample our random numbers¹:

sampleS :: RandomInt
sampleS =
  -- get the current state of the generator
  get                                 -- :: State Generator Generator
    >>= (\(_, gen) ->
      -- use it to generate a new random number
      let (num, gen') = sample gen in
      -- install the updated generator
      put gen'                        -- :: State Generator Generator
        -- install the generated number
        >>= (\_ -> num) )             -- :: State Generator Int

Then we can write simple wrapper functions to compute bare numbers:

-- in order to actually evaluate to a number, we need this function to evaluate
-- the state processor function
sampleInt :: Int -> Int
sampleInt seed = num
  where
    -- get the bare state processor function
    State proc = sampleS     -- :: Generator -> (Generator, Int)
    -- generate an initial state
    gen = makeGenerator seed -- :: Generator
    -- call the processor on the initial state, keeping only the output value
    (_, num) = proc gen      -- :: (Generator, Int)

-- just one number isn't so useful, so let's generate a list of numbers
intStream :: Int -> Int -> [Int]
intStream length seed = nums
  where
    -- this amounts to mapping each number in [1 .. length] to a random number
    items = [1 .. length]                     -- :: [Int]
    -- create a state processor that produces a list of output values
    randStream = mapM (\_ -> sampleInt) items -- :: State Generator [Int]
    -- unwrap the processor
    State proc = randStream                   -- :: Generator -> (Generator, [Int])
    -- generate an initial state
    gen = makeGenerator seed                  -- :: Generator
    -- call the processor on the initial state, keeping only the output value
    (_, nums) = proc gen                      -- :: (Generator, [Int])

-- `mapM` applies a monad-generating function to a list of items sequentially,
-- passing monadic output to the next item's function application with `>>=` and
-- accumulating unwrapped output in the monad
mapM :: Monad m => (a -> m b) -> [a] -> m [b]
mapM f items = wrapped
  where
    -- use `List` as a functor
    mapped = fmap f items -- :: [m b]
    -- compose all items in `mapped` using `>>=`, storing intermediate products
    wrapped = concatM (return []) mapped

concatM :: Monad m -> [m b] -> m [b]
concatM []            = return []
concatM (head : rest) =
  -- access the internal value of the first item
  head >>= (\val ->              -- head :: m b
    -- evaluate the action on the rest of the list
    (concatM rest) >>= (\vals -> -- rest :: [m b]
                                 -- concatM rest :: m [b]
      -- perform the concatenation of inner values
      return (val : vals) ) )    -- val : vals :: [b]

Footnotes

1. ↵ For a more concise implementation, check out do notation.

   sampleS :: RandomInt
   sampleS = do
     gen <- get
     let (num, gen') = sample gen
     put gen'
     return num