Understanding how reader monad simulates global variables

Let’s go through your example step by step in two ways. First, let’s see what the do notation equivalent is:

do
  a <- (+ 2)
  b <- (+ 3)
  return (a + b)

This is essentially an abstraction over the following repetitive pattern, where many different functions are all applied to the same shared “environment” value and then there’s some final result returned:

\ env -> let
  a = (+ 2) env
  b = (+ 3) env
  in (a + b)

In fact, for consistency with the other actions, the return also receives the environment, and just ignores it:

\ env -> let
  a = (+ 2) env
  b = (+ 3) env
  in const (a + b) env  -- const x y = x

We can make this fit the Monad pattern by having >>= carry this environment along and pass it to each function ((+ 2), (+ 3), const (a + b)).

>>= takes two functions function and continue, where function is one of (+ 2) or (+ 3) above; >>= applies function to the environment, and then passes two things along to continue: the result of function, and the environment again, so that subsequent steps in continue can also read the environment.

-- Type:
(>>=) :: (Monad m) => m        a -> (a -> m        b) -> m        b
-- or:
(>>=) ::              (->) env a -> (a -> (->) env b) -> (->) env b
-- or:
(>>=) ::              (env -> a) -> (a -> env ->   b) -> env ->   b

-- Implementation:
instance Monad ((->) env) where

  function >>= continue = \ env -> continue (function env) env
  -- or:
  (>>=) = \ function continue env -> continue (function env) env

In fact, with the {-# LANGUAGE InstanceSigs #-} extension, you can write exactly the last type above as the signature for >>= in the Monad instance, and I recommend using that extension to help you implement typeclasses as you’re learning.

return takes a value x, and produces a function that fits in place of continue above, which, as I mentioned before, just ignores the environment and returns x.

return x = \ _env -> x
-- =
return = \ x _env -> x
-- =
return = const

Now secondly, if we go back to the desugared do notation, and add parentheses to clarify the scope of the lambdas:

(+ 2) >>= (\ a -> (+ 3) >>= (\ b -> return (a + b)))

Then we can inline those definitions of >>= and return to see how they work in this context. I’ve numbered the variables in each substitution to make the scopes a little more explicit.

(\ function1 continue1 env1 -> continue1 (function1 env1) env1)
  (+ 2)
  (\ a ->
    (\ function2 continue2 env2 -> continue2 (function2 env2) env2)
      (+ 3)
      (\ b ->
        (\ x env3 -> x)
          (a + b)))

Then we can substitute all the variables throughout, by applying the functions to their arguments (beta-reduction):

let
  function1 = (+ 2)
  continue1 = \ a ->
    let
      function2 = (+ 3)
      continue2 = \ b ->
        let
          x = (a + b)
        in \ _env3 -> x

    in \ env2 -> continue2 (function2 env2) env2

in \ env1 -> continue1 (function1 env1) env1)

-- =

\ env1 ->
  (\ a env2 ->
    (\ b _env3 -> a + b)
      (env2 + 3)
      env2)
    (env1 + 2)
    env1

-- =

\ env1 -> let
  a = (env1 + 2)
  env2 = env1
  in let
    b = (env2 + 3)
    _env3 = env2
    in (a + b)

Then we can omit all the redundant env variables, since they’re all equal to each other, and just used to pass the value along to all the functions:

\ env -> let
  a = (env + 2)
  b = (env + 3)
  in (a + b)

And this is exactly the code we were trying to abstract away!

This is only one Monad instance, though; all other instances work differently based on the specific type involved. What all the basic instances have in common is that they abstract over some common repetitive pattern in pure code, e.g.:

-- State:

\ state1 -> let
  (a, state2) = function1 state1
  (b, state3) = function2 state2
  (c, state4) = function3 state3
  …
  in (x, stateX)

-- Writer:

\ log1 -> let
  (a, log2) = function1
  (b, log3) = function2
  (c, log4) = function3
  …
  in (x, log1 <> log2 <> log3 <> log4)

-- Maybe:

case function1 of
  Just a -> case function2 of
    Just b -> case function3 of
      Just c -> …
          Just x -> Just x
          Nothing -> Nothing
      Nothing -> Nothing
    Nothing -> Nothing
  Nothing -> Nothing

-- Either:

case function1 of
  Right a -> case function2 of
    Right b -> case function3 of
      Right c -> …
          Right x -> Right x
          Left errorX -> Left errorX
      Left error3 -> Left error3
    Left error2 -> Left error2
  Left error1 -> Left error1

-- List:

concatMap
  (\ a -> concatMap
    (\ b -> concatMap
      (\ c -> concatMap
          …
         [x])
      function3)
    function2)
  function1

Likewise, the basic monad transformers abstract over such patterns in monadic code.

-- MaybeT:

do
  ma <- function1
  case ma of
    Just a -> do
      mb <- function2
      case mb of
        Just b -> do
          mc <- function3
          case mc of
            Just c -> …
                Just x -> pure (Just x)
                Nothing -> pure Nothing
            Nothing -> pure Nothing
        Nothing -> pure Nothing
      Nothing -> pure Nothing
    Nothing -> pure Nothing

The Monad/Applicative/Functor hierarchy provide a common way for data structures and control structures to abstract over some kind of sequential pattern like this so they can all be used with do notation and library functions like replicateM that work for any Monad.

It helps to look at lots of different examples, not just one instance, to get a feel for what they have in common and how they differ.

In your question, you are describing one specific monad, the "reader" monad. At the level of detail you're trying to understand it, this reader monad will behave differently from other monads.

That said, I think your understanding of the detailed operation of the reader monad is partially correct. However, you've missed a couple of things: you've messed up the precise order of operations in your example expression that's dictated by how lambda expressions are parsed, and you've sort of glossed over the "merge" step that combines the \w -> ... lambdas that are introduced by each >>= operator so they all refer to the same global variable.

To make it a little clearer, let me spell it out in detail, using simple transformations of expressions that match the way GHC would actually evaluate your code...

Let's take your example:

(+2) >>= \a -> (+3) >>= \b -> return (a+b)

and start by inserting some parenthesis. Even though the operator >>= is left associative, lambda expressions like \a -> ... suck up as much valid expression as they can, so the correct placement of parentheses in this expression is:

(+2) >>= ( \a -> ( (+3) >>= ( \b -> return (a+b) ) ) )         -- (1)

Now, as you note, the definition of >>= is:

f >>= g = \w -> g (f w) w

Also, the definition of return is:

return x = \w -> x

so let's work on the expression for the right-most occurrence of >>= in (1):

(+3) >>= ( \b -> return (a+b) )
=   -- apply definition of `>>=` with f = (+3); g = \b -> return (a+b)
\w -> (\b -> return (a+b)) ((+3) w) w
=   -- apply lambda expression (\b -> ...) to first argument ((+3) w)
\w -> (return (a + (+3) w)) w
=   -- apply definition of `return` (with renamed variable w')
\w -> (\w' -> a + (+3) w) w
=   -- apply lambda expression (\w' -> ...) to first argument w
\w -> a + (+3) w

So, we have:

(+3) >>= ( \b -> return (a+b) )    ===    \w -> a + (+3) w     -- (2)

and, in a sense, the operator >>= works by passing the function (+3) as the argument b. However, it doesn't literally do this, because the argument b has numeric type, not function type. Instead, it "pulls out" a lambda for the "global variable" \w -> ... and then passes (+3) w as b, which works fine, since (+3) w has a numeric type.

Now, let's look at the original example (1) with the substitution (2):

-- from (1) and the subsitution (2):
(+2) >>= ( \a -> ( \w -> a + (+3) w ) )
=   -- apply defn of >>= with f = (+2); g = \a -> (\w -> a + (+3) w)
\w' -> (\a -> (\w -> a + (+3) w)) ((+2) w') w'
=   -- apply lambda (\a -> ...) to ((+2) w')
\w' -> (\w -> (+2) w' + (+3) w) w'
=   -- apply lambda (\w -> ...) to w'
\w' -> (+2) w' + (+3) w'

and, again, in a sense, the operator >>= works by passing the function (+2) as argument a. Since it can't do it literatlly, it does it by pullling a lambda for the "global variable" \w' -> ... out and passing (+2) w' as a, and then it sort of "merges" the \w -> ... lambda pulled out by the rightmost >>= operator with the \w' -> ... it's added, so they point to the same global variable.

Obviously, if we added additional similar binds on the left:

(+8) >>= \e -> (+1) >>= \o -> (+2) >>= \a -> (+3) >>= \b -> return (e+o+a+b)

we'd get the same pattern of adding a new lambda \w' -> ... for the global variable and then "merging" it with the existing lambda \w -> ...:

\w -> (+8) w + (+1) w + (+2) w + (+3) w