# You Could Have Invented The State Monad

I'm attempting NICTA/course a second time. I gave up the last time because none of the State exercises were making sense and I found myself leaning so heavily on the solutions that I wasn't actually learning anything. This time I was much better prepared after watching lots of CanFPG talks, reading lots of blog posts and writing a little Haskell, and I easily cleared the State hurdle. In fact, I'm now going to demonstrate how you (yes, you) could have come up with it (with a little help).

The fundamental insight of state is that it can be represented by a function that takes a value of type `s` and returns a tuple of some value `a` and a new value of type `s`:

```newtype State s a = State { runState :: s -> (a,s) }
```

Given such a type, what would its `Functor` instance look like?

```instance Functor (State s) where
(<\$>) :: (a -> b) -> State s a -> State s b
```

Our implementation should be another State that takes a value `s0`, passes it to the second argument `sa` (resulting in `(a, s1)`) and calls the function `fn` on `a`:

```  (<\$>) fn (State sa) = State (\s0 -> let (a, s1) = sa s0 in (fn a, s1))
```

This is a State that takes `s0` and returns `(b, s1)`, which is exactly what we wanted.

Let's look at the `Applicative` instance next:

```instance Applicative (State s) where
pure  :: a -> State s a
(<*>) :: State s (a -> b) -> State s a -> State s b
```

The implementation for `pure` explains where the `a` in our State comes from. Given some `a`, return a State that, when fed a value `s`, results in `(a,s)`. It practically writes itself.

```  pure a = State (\s -> (a,s))
```

`(<*>)` is a bit trickier, because we're dealing with both the State the function is in and the State its argument is in. The implementation should be a State that takes a value `s0`, feeds it to `sa` to get `(fn, s1)`, feeds `s1` to `sb` to get `(a, s2)`, and calls `fn` on `a`:

```  (<*>) (State sa) (State sb) =
State (\s0 -> let (fn, s1) = sa s0
(a,  s2) = sb s1
in (fn a, s2))
```

The hardest thing is remembering to thread `s0` through `sa` and `sb` so that we don't lose any state on the way. We can usually follow the types but they don't help in this specific case.

Finally, let's look at the `Monad` instance:

```instance Monad (State s) where
(>>=) :: State s a -> (a -> State s b) -> State s b
```

As with all our previous implementations, it has the form:

```  (>>=) (State sa) fn = State (\s0 -> let ??? in ???)
```

We know that we need to feed `s0` to `sa` to get an `a` to apply to `fn`:

```  (>>=) (State sa) fn =
State (\s0 -> let (a, s1) = sa s0
???     = fn a
in ???)
```

The result of `fn a` is a `State sb` but we need to return a tuple of `(b, s)`. We can obtain one by feeding `s1` to `sb`:

```  (>>=) (State sa) fn =
State (\s0 -> let (a, s1)  = sa s0
State sb = fn a
in sb s1)
```

Success!

Let's define a few functions to make our lives easier. `get` returns a State that, when fed some `s`, returns `(s,s)`. This allows us to expose `s` for direct modification:

```get :: State s s
get = State (\s -> (s, s))
```

`put` allows us to store a State that ignores the `s` passed to it later:

```put :: s -> State s ()
put s = State (\_ -> ((),s))
```

Sometimes we want the `s` and not the `a`:

```exec :: State s a -> s -> s
exec (State sa) s = snd \$ sa s
```

At other times we want the `a` and not the `s`:

```eval :: State s a -> s -> a
eval (State sa) s = fst \$ sa s
```

With all this machinery in place, we can do this:

```Prelude> exec (do i <- get; put (i+1); return ()) 0
1
```

I still couldn't believe that this worked the first time I tried it, so let's desugar this:

```    do i <- get; put (i+1); return ()
== get >>= \i -> put (i+1) >>= \_ -> pure ()
== State (\s -> (s, s))    >>= \i ->
State (\_ -> ((), i+1)) >>= \_ ->
State (\s -> ((), s))
```

Let's simplify from the bottom up. By the definition of `(>>=)`:

```    (>>=) (State (\_ -> ((), i+1))) (\_ -> (State (\s -> ((), s)))) =
State (\s0 -> let (a, s1)  = (\_ -> ((), i+1)) s0
-- (a, s1)  = ((), i+1)
State sb = (\_ -> (State (\s -> ((), s)))) a
--       sb = (\s -> ((), s))
in sb s1)
-- ((), i+1)
== State (\s0 -> ((), i+1))
== State (\_  -> ((), i+1))
```

Plugging that back in, we have

```State (\s -> (s,s)) >>= \i -> State (\_ -> ((), i+1))
```

Which we can simplify in the same way:

```    (>>=) (State (\s -> (s,s))) (\i -> State (\_ -> ((), i+1))) =
State (\s0 -> let (a, s1)  = (\s -> (s,s)) s0
-- (a, s1)  = (s0, s0)
State sb = (\i -> State (\_ -> ((), i+1))) a
--       sb = (\_ -> ((), s0+1))
in sb s1)
-- ((), s0+1)
== State (\s0 -> ((), s0+1))
== State (\i  -> ((), i+1))
```

Finally, we have

```    exec (State (\i -> ((), i+1))) 0
== snd \$ runState (State (\i -> ((), i+1))) 0
== snd \$ (\i -> ((), i+1)) 0
== snd \$ ((), 1)
== 1
```

This is my favourite thing about Haskell: the fact that it is built on abstractions that can be reasoned about in such a rigorous manner.

In fact, with some inspired renaming, you too could have invented the continuation monad.