Deriving the State Monad

I am currently learning Haskell. When I encountered Monads, I wanted to understand it from examples and first principles.

This article is my attempt to understand Monads, particularly the State Monad purely from examples.

The first video about Monads which made a lot of sense was this video from Bartosz Milewski’s category theory course on youtube:

In which defines an operator called fish

>=> :: (a → m b) -> (b → m c) → (a → m c)
-- or lets call it as fish:
fish :: (a → m b) -> (b → m c) → (a → m c)
fish f g = \a = let mb = f a in
bind mb g
bind :: m b -> (b -> m c) -> m c

The function fish is defined using another function bind which helps it unpack the parameter of type m b to b and apply function g to get the return value of type m c

bind can be defined as of type: m a -> (a -> m b) -> m c

Thus bind helps in composing functions of type a-> m b where m b is a Haskell type.

With this background I set to define my own Monads:

I first derived a Monad for Maybe as follows:

data Maybe' a = Nothing' | Just' a deriving Show
sqrt' :: (Floating a, Ord a) => a -> Maybe' a
sqrt' x = if x < 0 then Nothing' else Just' (sqrt x)
inv' :: (Floating a, Ord a) => a -> Maybe' a
inv' x = if x == 0 then Nothing' else Just' (1/x)
log' :: (Floating a, Ord a) => a -> Maybe' a
log' x = if x == 0 then Nothing' else Just' (log x)

A function that composes sqrt’, inv’ and log’ would look like:

sqrtInvLog' :: (Floating a, Ord a) => a -> Maybe' a
sqrtInvLog' x = case (sqrt' x) of
Nothing' -> Nothing'
(Just' y) -> case (inv' y) of
Nothing' -> Nothing'
(Just' z) -> log' z

This function could be simplified by factoring the common code pattern of

case x of
Nothing -> Nothing
(Just y) -> f y

into a function:

fMaybe' :: (Maybe' a) -> (a -> Maybe' b) -> Maybe' b
fMaybe' Nothing' _ = Nothing'
fMaybe' (Just' x) f = f x

Using this function sqrtInvLog' becomes simpler:

-- Applying fMaybe' =>
sqrtInvLog' :: (Floating a, Ord a) => a -> Maybe' a
sqrtInvLog' x = (sqrt' x) `fMaybe'` (inv') `fMaybe'` (log')

Now we can generalise this concept using Monads:

class Monad' m where
bind' :: m a -> (a -> m b) -> m b
return' :: a -> m a
instance Monad' Maybe' where
bind' Nothing' _ = Nothing'
bind' (Just' x) f = f x
return' x = Just' x

Using Monad’ ,sqrtInvLog' now can be written as:

sqrtInvLog' :: (Floating a, Ord a) => a -> Maybe' a
sqrtInvLog' x = (sqrt' x) `bind'` (inv') `bind'` (log')

Next step is to see how can apply this concept to maintain state. For example we could maintain state which is the sum of all intermediate step of sqrt, inv and log.

Thus the return value of sqrtInvLog would be the end result by applying sqrt, inv and log serially and state would be thesum of all intermediate values.

My first attempt was as follows:

data St a s = St (a,s) deriving Show
sqrtLogInvSt :: (Floating a, Ord a) => St a a -> St (Maybe' a) a
sqrtLogInvSt (St (x,s)) = case (sqrt' x) of
Nothing' -> St (Nothing', s)
(Just' y) -> case (log' y) of
Nothing' -> St (Nothing', s+y)
(Just' z) -> St (inv' z, s+y+z)

The trouble with this definition of stats as St is that it is not possible to define a Monad using the above definition as bind needs to be defined as taking a single type parameter, where asSt takes a and s

This is where I hit a roadblock and hit Stackoverflow. Answers to my this question indicated a hint about trying to using the following Haskell definition of State Monad :

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

Given this, I tried to follow the same pattern I had used for Maybe. The pattern of first attempting to define a composed function, from which I can then extract a bind like function (similar to fMaybe') :

-- A sample function that returns a State
fex1 :: Int->State Int Int
fex1 x = State { runState = \s->(r,(s+r)) } where r = x `mod` 2
-- Another
fex2 :: Int->State Int Int
fex2 x = State { runState = \s-> (r,s+r)} where r = x * 5

My attempt at a composed function looked like this (with zero as the initial state):

fex3 :: Int-> (Int, Int)
fex3 x = (runState (fex2 y)) st where (st, y) = (runState (fex1 x)) 0

This definition was a mistake which led me down a wrong path. The mistake was that the definition should have been similar to sqrtIngLog and returned type State Int Int but instead it hardcoded the initial state and returned type (Int, Int)

But still going ahead I attempted to define a Monad.

State as well takes two type parameters. But I attempted to define the Monad as follows:

instance Monad’ (State s) where
bind' st f = undefined
return' x = State { runState = \s -> (x,s) }

And this complies. The important point (which I later realised) is that

(State s) is a partially applied type.

At this point I was stuck unable to come up with a definition for bind. I had to hit stackoverflow again, where this answer resolved the puzzle:

The answer corrected my mistake of wrong type for the composed function. It should have been as follows:

fex3 :: Int -> State Int Int
fex3 x = State { runState = \s ->
let (y, s') = (runState (fex1 x)) s in
runState (fex2 y) s'

This mirrors the composed function of sqrtInvLog .

Taking the composition further, we can define a function that composes three functions:

fex4 :: Int -> State Int Int
fex4 x = State { runState = \s ->
let (y, s') = (runState (fex1 x)) s in
let (y', s'') = (runState (fex2 y)) s' in runState (fex2 y') s''

The common pattern should apply a function f given an earlier state output (from an earlier function) aka fMaybe'like is:

(sqrt’ x) `fMaybe’` (inv’) `fMaybe’` (log’)`

Mirroring the above definition for fex4 it would be

(fex1 x) `bindState` (fex2) `bindState` (fex2)

bindState can be implemented as:

bindState st f = State { runState = \s ->
let (x, s') = (runState st) s in
(runState (f x)) s'
}

Which now can be generalised as the State Monad:

instance Monad' (State s) where
bind' st f = State { runState = \s ->
let (y, s') = (runState st) s in
(runState (f y)) s'
}
return' x = State { runState = \s -> (x,s) }

fex4 can now be written as:

fex4' x = (fex1 x) `bind'` (fex2) `bind'` (fex2)

Another important point is that State s a takes two type parameters and does not fit the pattern of m a; but a partially applied type (State s)does fit (m a) (as State s a can be interpreted as (State s) a which does fit m a . Thus allowing us to define

instance Monad’ (State s) where
bind’ st f = ...