Defining the behaviour of >>= for a monad

Question

I'm learning about monads in Haskell at uni and this is an exercise we've been given:

Given the type

data Expr a = Var a | Val Int | Add (Expr a) (Expr a)

of expressions built from variables of type a, show that this type is monadic by completing the following declaration:

instance Monad Expr where
--  return         :: a -> Expr a
    return          = ...

--  (>>=)          :: Expr a -> (a -> Expr b) -> Expr b
    (Var a)   >>= f = ...
    (Val n)   >>= f = ...
    (Add x y) >>= f = ...

So far I've logically managed to come up with:

return        = Var

(Var a) >>= f = f a
(Val n) >>= f = Val n

But how do you define (Add x y) >>= f?

Answer 1

This is a case where we can just "follow our noses" by examining the types: when given Add x y we have two things of type Expr a (x and y) and we're trying to output something of type Expr b. One option would be to just entirely ignore x and y and output something like Val 0 (but that's not very interesting, and probably breaks the Monad laws).

Hence, assuming we wish to use both x and y a good first step would be to transform those two Expr as into Expr bs: we have exactly one choice here since there's exactly one relationship between a and b, the function f :: a -> Expr b.

Now, to use these two pieces of information, we can think about looking for a function with type Expr a -> (a -> Expr b) -> Expr b, that is, something to take our expression (Expr a) and the function (a -> Expr b) and give us an Expr b.

Seem familiar?

It should: it's exactly what we're working with, the bind operator >>=. By using this, we've now got two Expr bs: (x >>= f) and (y >>= f).

Hence, the last step is to use these two things to give the result: a single Expr b... there's a single choice for this too.

Answer 2

This is a game of typing tetris. Let's consider the types of your first two (correct) branches.

a :: a
f :: a -> Expr b
----------------
Expr b

We read this as "given all of things above the line, we must produce something of the type listed below the line". In this case, the answer is obvious: we just apply f to a.

a :: a
f :: a -> Expr b
----------------
f a :: Expr b

---

For Val n we can repeat the process

n :: Int
f :: a -> Expr b
----------------
Expr b

at first this seems impossible, but we have elided a few things "above the line" that are actually available. Importantly, there's this:

n :: Int
f :: a -> Expr b
Val :: forall x . Int -> Expr x
--------------------
Expr b

and since x is not specific (it's bound by forall) it can unify with b.

n :: Int
f :: a -> Expr b
Val :: forall x . Int -> Expr x
--------------------
Val n :: Expr b

---

For the final case, we have

x   :: Expr a
y   :: Expr a
f   :: a -> Expr b
Add :: forall x . Expr x -> Expr x -> Expr x
--------------------------------------------
Expr b

Again, this seems impossible. We could reassemble the Add using x and y but that just gives us a type Expr a instead of Expr b. We can't apply the f because it needs an a not an Expr a.

So the trick is that with recursive data types you're almost certainly going to use recursive definitions to functions... So let's bring in one other thing we have from the environment.

x     :: Expr a
y     :: Expr a
f     :: a -> Expr b
Add   :: forall x . Expr x -> Expr x -> Expr x
(>>=) :: forall y z . Expr y -> (y -> Expr z) -> Expr z
-------------------------------------------------------
Expr b

Again, due to the forall we can use (>>=) on Expr types whatever their variable is. Almost immediately we have only one way to go forward: the only value we have that could be the second argument to (>>=) is f

x       :: Expr a
y       :: Expr a
Add     :: forall x . Expr x -> Expr x -> Expr x
(>>= f) :: Expr a -> (a -> Expr b) -> Expr b
-------------------------------------------------------
Expr b

And now we could apply either x or y on the left side of (>>= f) to get a value of type Expr b. Unfortunately, both of these are wrong. One way to be sure of this is that for highly general functions like (>>=) we almost never throw away information—each argument should be used non-trivially.

Fortunately, if we have two Expr bs we can use Add to combine them:

x       :: Expr a
y       :: Expr a
Add     :: Expr b -> Expr b -> Expr b
(>>= f) :: Expr a -> (a -> Expr b) -> Expr b
-------------------------------------------------------
Expr b

and now we have a way to use both x and y non-trivially

x       :: Expr a
y       :: Expr a
Add     :: Expr b -> Expr b -> Expr b
(>>= f) :: Expr a -> (a -> Expr b) -> Expr b
-------------------------------------------------------
Add (x >>= f) (y >>= f) :: Expr b

noting that we keep the x and y in the same order.

---

So this while game of type-tetris is a bit long-winded but demonstrates that with a few principles:

1. Definitions over recursive types tend to need recursion
2. Definitions of general functions will use every argument non-trivially if possible

we can get to the answer to how to define (>>=) almost entirely mechanically. A good definition ought to feel satisfying, like there literally is no other choice.

instance Monad Expr where
  Var a >>= f = f a
  Val n >>= _ = Val n
  Add x y >>= f = Add (x >>= f) (y >>= f)
  ...

We can interpret this definition as "reacting" the f and the lead nodes of Var where as exist. If no a exists then we're free to do nothing. If recursive Exprs exist then we just "push the call to (>>= f) down" and rebuild the recursive type.