Hackle's blog
between the abstractions we want and the abstractions we get.

# State Monad: a bit of currying goes a long way

A friend asked me about the `State` Monad. He enjoyed a honey-moon with Monads such as `Maybe`, `List` and `Either`, but `State` appears to be next-level. It brings non-local states, it encodes a function, and often appears to be dark magic.

So I tried to find him a good introduction, and got myself mighty confused in the process - such is the curse of Monad tutorials!

But `State` is nothing more than good ol' currying and partial application with a flourish. So here comes a very special introduction for those who are already comfortable with the idea of Monads, are not scared of type signatures, and don't like beating around the bush with motivational examples.

This leaves me an audience of about 10 people, so I'll get right to it.

## (State s a) type by type

``````-- take a simple function
(a -> b)
-- we can add state s to the side to simulate a "stateful" computation
(a, s) -> (b, s)

-- what about function composition? (Note: this is pre-composition for convenience)
(a -> b) -> (b -> c) -> (a -> c)
-- do the same thing: add state s, and apply a few transformations
((a, s) -> (b, s))   -> ((b, s) -> (c, s))   -> ((a, s) -> (c, s))
== (a -> s -> (b, s))   -> (b -> s -> (c, s))   -> (a -> s -> (c, s))    -- currying
== (a -> (s -> (b, s))) -> (b -> (s -> (c, s))) -> (a -> (s -> (c, s)))  -- partial application
newtype State s a = State (s -> (a, s)) -- this is the "stunt"
== (a -> State s b)     ->  (b -> State s c)    -> (a -> State s c)      -- substitution
== (a -> m b) -> (b -> m c) -> (a -> m c)
== (>=>)    -- "Kleisli arrow" from Control.Monad
``````

Note this is not a proof, but an intuition built with transformation of types step-by-step, from function composition to Monad composition.

Finding this a bit dense? No worries, I've got some examples anyway.

## the JavaScript envy

Imagine a stateful function in a language more "free" with mutation, such as JavaScript

``````const greetings = [ "Hello", "Howdy", "Hi", "G'day" ];

// this is the stateful part
let index = 0;

function greet(name: string): string {
return greetings[index++];  // I am leaving out boundary check for clarity
}

console.log(greet("Hackle"));
> Hello Hackle
console.log(greet("Hackle"));
> Howdy Hackle
``````

See how the state `index` is messing with the `greet` function, so it's indeterministic?

It we try to translate it to Haskell,

``````greetings = [ "Hello", "Howdy", "Hi", "G'day" ]
index = 0

type Name = String
type Greeting = (String, Name) -- why not a string? See below.

greet :: Name -> Greeting
greet name = (greetings !! index, name)
``````

(A small note, I use `(String, Name)` instead of concatenating to a single `String`, to avoid getting two `String`s mixed up in the steps to come.)

Just one problem - there is no easy way to mutate `index`. So a direct translation from JavaScript is a no-go. What can we do?

### Keep the state on the side

Let's try again in the Haskell way. We are going to keep the "state" on the side, as, passing it around as an extra parameter.

``````-- I am going to call it "state" now to suit the narrative :-)
type GState = Int

greetS :: (Name, GState) -> (Greeting, GState)
greetS (name, st) = ((greetings !! st, name), st + 1)

ghci> greetS ("Hackle", 0)
(("Hello","Hackle"),1)

ghci> let (greeting, st) = greetS ("Hackle", 0) in [greeting, fst \$ greetS("Hackle", st)]
[("Hello","Hackle"),("Howdy","Hackle")]
``````

This simulates a stateful function, but syntax-wise, it's annoying to thread the state through, and just not as sweet as JavaScript! Can we do better?

### Composition

Some of my friends do what I call "double greeting" - instead of simply saying "hello Hackle!", they go "Hi Hackle, G'day!". Let's model that with function `doubleGreetS`, which turns a `Greeting` into a `DoubleGreeting`. (Again, I use a 3-tuple instead of a plain `String` to avoid confusion).

``````type DoubleGreeting = (String, Name, String)

doubleGreetS :: (Greeting, GState) -> (DoubleGreeting, GState)
doubleGreetS ((greeting, name), st) = ((greeting, name, greetings !! st), st + 1)

ghci> doubleGreetS (("Hello", "Hackle"), 1)
(("Hello","Hackle","Howdy"),2)
``````

We see the input to `doubleGreetS` matches the output of `greetS`, so let's create a `composeS` that stitches them together.

``````composeS
:: ((Name, GState) -> (Greeting, GState))
-> ((Greeting, GState) -> (DoubleGreeting, GState))
-> ((Name, GState) -> (DoubleGreeting, GState))
composeS f1 f2 =
\(a, st) ->
let (b, st1) = f1 (a, st)
in f2 (b, st1)

highGreet :: (Name, GState) -> (DoubleGreeting, GState)
highGreet = composeS greetS doubleGreetS

ghci> highGreet (("Hackle"), 0)
(("Hello","Hackle","Howdy"),2)
``````

Experienced Haskellers would immediately point out `composeS` is just `.` in disguise! Keep this in mind. We will continue to work on the more verbose `composeS` signature, but the aim is to arrive at something comparable to `.`.

### Stunt 1: currying!

Watch out - I am going to pull a stunt!

``````composeS
::  ((Name, GState)     -> (Greeting, GState))
->  ((Greeting, GState) -> (DoubleGreeting, GState))
->  ((Name, GState)     -> (DoubleGreeting, GState))

-- currying
(Name       -> GState -> (Greeting, GState))
->  (Greeting   -> GState -> (DoubleGreeting, GState))
->  (Name       -> GState -> (DoubleGreeting, GState))

-- partial application, notice the pattern?
(Name       -> (GState -> (Greeting, GState)))
->  (Greeting   -> (GState -> (DoubleGreeting, GState)))
->  (Name       -> (GState -> (DoubleGreeting, GState)))
``````

Oh well, I exaggerate. It's simply a routine transformation: currying + partial application. And by "partial application" I really just mean adding `()`.

(Also I am leaving out the updates to the implementation of `composeS` as an exercise for the reader.)

As simple as the "stunt" is, it does surface something interesting - A pattern is winking at us, if we quint a bit: `a -> (s -> (b, s))`. Try to find it yourself!

With this finding, we can generalise the last type to,

``````    (a -> (s -> (b, s)))
->  (b -> (s -> (c, s)))
->  (a -> (s -> (c, s)))
``````

Remember `s` stands for "State"? Looking at this re-organised type, we will further notice the repetition in `s -> (?, s)`. Indeed, this is the key to our topic at hand.

### The State Monad: an underwhelming introduction

As developers do, when there is repetition, we create a type. In this case, smart people figured out that we can create the famous `State` type. Behold!

``````newtype State s a = State (s -> (a, s))
``````

(Note `a` is polymorphic - so is `s` for that matter - it can be any type, `String`, `Int`, or `b`, `c`).

With the tedious lead-up, this may appear underwhelming. But if you have tried other introductions that start with this type, it's fair to say, `s -> (a, s)` is not made for fast digestion.

Let's hold the celebration just yet. Try answer this: what's the intuition for `s -> (a, s)`?

A naive interpretation is, a value of type `a` can be computed from state `s`, like turning `String` to `Int` with `read`. While this can be the case for some use of `State`, it's not always true, and does not necessarily have to be so.

The more sophisticated interpretation, is the function `s -> (a, s)` has the "knowledge" of producing an `a`. How is that possible? Why, I am surprised you'd ask! One way is partial application, which we've just seen so much of!

In our example, `a -> s -> (b, s)` can be partially applied with `a`, to leave us `s -> (b, s)`. Ta da.

To continue the work on the types, let's slot in `State s a` (again, `a` is polymorphic so it can be `b` or `c`!)

``````    (a -> State s b)
->  (b -> State s c)
->  (a -> State s c)
``````

Already much easier on the eye, wouldn't you say? At least we saved 2 layers of `()`.

If you try to catch up with the implementation, there would be a fair bit of wrapping and unwrapping with `State` being `newtype`; but if we focus only on the composition, it should remind us of monad composition. Presumed we can prove `State` is a monad, the above types can be generalised to,

``````    (a -> m b)
->  (b -> m c)
->  (a -> m c)
``````

Hello, it's none other than the fish operator `>=>`, or the "Kleisli arrow"! According to hoogle.

Before it's too late, we need to implement the monad type class. Luckily this is straightforward (try it out yourself!). Below is a very naive version.

``````instance Functor (State s) where
fmap a2b (State f) = State \$ \s -> let (a, s1) = f s in (a2b a, s1)

instance Applicative (State s) where
pure a = State \$ \s -> (a, s)
(State g) <*> (State f) = State \$ \s -> let (a, s1) = f s; (a2b, s2) = g s1 in (a2b a, s2)

return = pure
(State f) >>= g = State \$ \s -> let (a, s1) = f s; (State h) = g a in h s1
``````

So `State s` is a proper monad, bravo! What does that mean for the example?

First, it prompts the refactoring below. I've suffixed the names with `M` to indicate the monad usage.

``````greetM :: Name -> State GState Greeting
greetM name = State \$ \st -> ((greetings !! st, name), st + 1)

doubleGreetM :: Greeting -> State GState DoubleGreeting
doubleGreetM (greeting, name) = State \$ \st -> ((greeting, name, greetings !! st), st + 1)

highGreetM = greetM >=> doubleGreetM

ghci> let (State f) = highGreetM "Hackle" in f 0
(("Hello","Hackle","Howdy"),2)
``````

You'll notice the immediate consequence of currying and partial application: what we used to supply in one go for `highGreet ("Hackle", 0)` is now done in two steps,

1. first, `"Hackle"` is given to the monad-powered `highGreetM`, which returns a `State`-monad value that encodes a partially-applied function `s -> (a, s)`, which
2. accepts the second parameter `0`, to complete the computation, and give us the same result as the non-monad-powered `highGreet`!

Despite the small win that the function types are more revealing by indicating state usage alongside return type, let's be honest, this consequence does not improve the life of the caller, and it's arguable if the implementation of `greetM` or `doubleGreetM` is any simpler. (I hear you, it's fun to use the "fish" operator). Not to forget, this is still a far cry from the "beauty" of the JavaScript code.

That's fair! I am not offended, because we aren't done yet! After all, how could "statefulness" be claimed without `putState` and `getState`?! Behold...

### Stunt 2: getter and setter

The famous `getState` is defined as,

``````getState :: State s s
getState = State \$ \s -> (s, s)
``````

You can see `State s s` is just a clever trick - if `State s a` can have any `a`, why not `s` to make it `State s s`?

Standalone, `getState` looks pretty silly. However, taken in the context of monad composition, it's nothing short of genius, because it allows us to grab the state out of thin air.

`putState` is reminiscent of the imperative "setter" that sets the state and returns `void`.

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

Not much is happening - it ignores any previous state with `_`, and sets up a void-like `()`. The only useful thing it does is putting an `s` in the second position of a tuple. This is key - have a look at the `>>=` definition, you'll see how this is enough to achieve the goal of "updating" the state (for any following `State`-monad values).

Now our example looks suspiciously familiar,

``````greetM' :: Name -> State GState Greeting
greetM' name = do
st <- getState
putState (st + 1)
return (greetings !! st, name)

doubleGreetM' :: Greeting -> State GState DoubleGreeting
doubleGreetM' (greeting, name) = do
st <- getState
putState (st + 1)
return (greeting, name, greetings !! st)
``````

Are you getting the imperative vibe? (I must admit seeing the `putState` and `return` "statements" does give me creeps!)

Compare this to `greetM`, see how state handling is offloaded to `getState` and `putState`, and separate from building the greeting? Pretty neat!

### the full imperative vibe

The imperative feel doesn't stop here. Consider this challenge: how do you map and sum a list of integers in one go?

One way to do it would be,

``````mapSum1 :: Num a => (a -> b) -> [a] -> (a, [b])
mapSum1 f xs = foldr (\x (tot, ys) -> (tot + x, f x : ys)) (0, []) xs
``````

Looks familiar? It's our old friend keep-the-state-on-the-side.

Even with this simple example, it's obvious that using a tuple is noisy. Now we have grokked `State`, it's time to introduce `mapM`: the "stateful" `map`.

``````mapSum2 f xs = mapM (\x -> do { tot <- getState; putState (tot + x); return (f x); }) xs

ghci> let (State f) = mapSum2 show [1..5] in f 0
(["1","2","3","4","5"],15)
``````

Of course we've come too far to not also introduce `modifyState`.

``````modifyState :: (s -> s) -> State s ()
modifyState f = do
s <- getState
putState (f s)

mapSum3 f xs = mapM (\x -> do { modifyState (+ x); return (f x); }) xs
ghci> let (State f) = mapSum3 show [1..5] in f 0
(["1","2","3","4","5"],15)
``````

Still with me? Congratulations, you can now write JavaScript in Haskell.

## Acknowledgement

Many thanks to Utku Demir for the comprehensive and on-point review.