In this post, I am going to introduce indexed monads, which generalize monads with additional type parameters carrying the information about the computational effects.

Motivation

The State monad represents computations with a state that can be queried and updated. For example, an Int state is queried and updated during the computation of c in the following example. While the value of the state is changed from 0 to 1, the type of the state remains the same during the entire computation.

import Control.Monad.State

test1 = runState c (0::Int) where
         c = do
             v <- get
             put (succ v)
             return v
-- (0, 1)

This is okay in most cases, but we sometimes want to express a computation where not only the value but also the type of the state can be changed. The vanilla State monad is not general enough to express this requirement.

Indexed Monads

Indexed monads are a generalization of monads that index each monadic type by an initial (type)state and a final (type)state. m is a type constructor for three type arguments, p, q and a. The argument a is the type of values produced by the monadic computation. p and q represent the types of the state before and after the computation.

class IxMonad m where
    ireturn :: a -> m p p a
    ibind :: m p q a -> (a -> m q r b) -> m p r b

ireturn and ibind must meet the monad laws as the ordinary monads do. ibind is required to be associative and ireturn to be the left and the right unit of ibind.

All ordinary monads can be injected into IxMonad with a newtype wrapper MW. It is a phantom type as the type parameters p and q are not used on the right hand-side of MW.

newtype MW m p q a = MW { unMW:: m a }

instance Monad m => IxMonad (MW m) where
    ireturn = MW . return
    ibind (MW m) f = MW (m >>= unMW . f)

Here is an example of using the ordinary State monad wrapped with MW. iget and iput wraps the result with MW newtype wrapper.

iget :: (MonadState s m) => MW m s s s
iget = MW get

iput :: (MonadState s m) => s -> MW m s s ()
iput = MW . put

test2 = runState (unMW c) (0::Int) where
         c = iget `ibind` (
               \v -> iput (succ v) `ibind` (
                 \_ -> ireturn v))
-- (0, 1)

Indexed State Monad

IxStateT defines an indexed state monad where si and so represents the input and the output state type respectively. The definition of IxStateT is similar to that of StateT except that the type of the state can be changed during the computation.

newtype IxStateT m si so v = IxStateT { runIxStateT:: si -> m (so,v) }

instance Monad m => IxMonad (IxStateT m) where
  ireturn x = IxStateT (\si -> return (si,x))
  ibind (IxStateT m) f = IxStateT (\si -> m si >>= (\ (sm,x) -> runIxStateT (f x) sm))

vsget :: Monad m => IxStateT m si si si
vsget = IxStateT (\si -> return (si,si))

vsput :: Monad m => so -> IxStateT m si so ()
vsput x = IxStateT (\si -> return (x,()))

The following example gets an Int from the state and puts a String into the state. We can see that the type of the state is changed from Int to String.

test3 = runIxStateT c (0::Int) >>= print where
         c = vsget `ibind` (
               \v -> vsput (show v) `ibind` (
                 \_ -> vsget `ibind` (
                   \v' -> ireturn (v,v'))))
-- ("0",(0,"0"))

Do notation

The IxMonad examples above looks ugly as we couldn’t use the do notation. Fortunately, -XRebindableSyntax extension allows us to overload the do-notation by providing alternative definitions that are local to the module.

{-# LANGUAGE RebindableSyntax #-}

import Prelude hiding ((>>=), (>>), return)
import IxState

return :: (Monad m) => a -> IxStateT m si si a
return = ireturn

(>>=) :: (Monad m) => IxStateT m p q a -> (a -> IxStateT m q r b) -> IxStateT m p r b
(>>=) = ibind

(>>) :: (Monad m) => IxStateT m p q a -> IxStateT m q r b -> IxStateT m p r b
v >> w = v >>= \_ -> w

c :: (Monad m) => IxStateT m Int String (Int, String)
c = do
  v <- vsget
  vsput (show v)
  v' <- vsget
  return (v, v')

Other definitions

There are multiple ways to define indexed monads. The one used here is from Robert Atkey’s Parameterised Notions of Computation.

Other definitions include:

• McBride: Kleisli Arrows of Outrageous Fortune
• Orchard: Fun with indexed monads

References

1. Oleg Kiselyov’s Parameterized `monad’
2. Indexed Monad section of Stephen Diehl’s What I Wish I Knew When Learning Haskell

Novice Haskell programmers think that monads are only for IO and stateful computations. But experienced Haskell programmers use monads to better structure their programs.

In this blog post, I am going to show you how we can better organize our code using monads.

Motivating example: Lambda lifting

Lambda lifting is a compiler transformation which eliminates all free variables from function definitions. It is an important step in a lazy functional language because it greatly simplifies evaluation on the graph reduction machine.

In his paper, Simon Peyton Jones describes how to perform lambda lifting in a modular fashion. The lambda lifter works in three steps:

-- | freeVars: Annotate every node in the expression with its free variables.
freeVars :: Expression -> AnnExpr Name (Set Name)

-- | Abstract the free variables from each lambda abstraction, replacing the lambda abstraction with the application of the new abstraction.
abstract :: AnnExpr Name (Set Name) -> Expression

-- | Give a unique name to each supercombinator and collect all the supercombinator definitions.
collectSCs :: Expression -> [SCDefn]

lambdaLift is the composition of these three functions.

lambdaLift :: Expression -> [SCDefn]
lambdaLift = collectSCs . abstract . freeVars

I am not going to explain the details of these steps in this post. Interested readers are referred to SPJ’s paper and book.

Instead, let’s dive into the last step and see how collectSCs is actually implemented.

Collecting supercombinators

collectSCs is defined in terms of a helper function named collecSC' which returns both the collection of supercombinators it has found and the transformed expression. It also carries around a name supply as an argument and returns the depleted supply as a result because it needs to generate fresh names for supercombinators.

-- | Gives a unique name to each supercombinator, collects all the
-- supercombinator definitions into a single list, and introduce the
-- $main supercombinator definition.
collectSCs :: Expression -> [SCDefn]
collectSCs e = ("$main", [], e') : scs
  where
  (_, scs, e') = collectSCs' initialNameSupply e

collectSCs' :: NameSupply -> Expression -> (NameSupply, [SCDefn], Expression)
collectSCs' ns (EConst k) = (ns, [], EConst k)
collectSCs' ns (EVar v) = (ns, [], EVar v)
collectSCs' ns (EAp e1 e2) =
  (ns2, scs1 ++ scs2, EAp e1' e2')
  where
  (ns1, scs1, e1') = collectSCs' ns e1
  (ns2, scs2, e2') = collectSCs' ns1 e2
collectSCs' ns (ELam args body) =
  (ns2, (name, args, body') : bodySCs, EConst (CFun name))
  where
  (ns1, bodySCs, body') = collectSCs' ns body
  (ns2, name) = newName ns1 "SC"
collectSCs' ns (ELet isRec defns body) =
  (ns2, scs, ELet isRec defns' body')
  where
  (ns1, bodySCs, body') = collectSCs' ns body
  ((ns2, scs), defns') = mapAccumL collectSCs'' (ns1, bodySCs) defns

  collectSCs'' (ns, scs) (name, rhs) =
    ((ns1, scs ++ scs'), (name, rhs'))
    where
    (ns1, scs', rhs') = collectSCs' ns rhs

The code is rather complex compared to what it actually does. The only place where interest things happen is lambda abstractions. It replaces lambda abstractions by names and return supercombinators.

The code is complex because it violates the most important software engineering principle: separation of concerns. collectSCs contains at least three orthogonal concerns:

1. Generation of fresh names
2. Accumulation of supercombinatros
3. Transformation of expressions

Organize code using monads

Monads are great tools to separate concerns. For example, Reader monad helps us get rid of an extra argument used to pass a context. Writer monad frees us from the agony of returning the accumulated results in every function.

So let’s separate our concerns in collectSCs.

1. Supply monad for fresh name generation (it’s a State monad in disguise)
2. Writer monad for accumulation of supercombinators

Because we need to compose two different monads, we use the monad transformer SupplyT.

type Collector a = SupplyT (Writer [SCDefn]) a

collectSCs :: Expression -> [SCDefn]
collectSCs e = ("$main", [], e') : scs
  where
  (e', scs) = runWriter $ evalSupplyT 0 (collectSCs' e)

collectSCs' :: Expression -> Collector Expression
collectSCs' (EConst k) = return $ EConst k
collectSCs' (EVar v) = return $ EVar v
collectSCs' (EAp e1 e2) = do
  e1' <- collectSCs' e1
  e2' <- collectSCs' e2
  return $ EAp e1' e2'
collectSCs' (ELam args body) = do
  name <- freshName
  body' <- collectSCs' body
  collect (name, args, body')
  return $ EConst (CFun name)
collectSCs' (ELet isRec defns body) = do
  body' <- collectSCs' body
  defns' <- traverse (\(name, defn) -> (name,) <$> collectSCs' defn) defns
  return $ ELet isRec defns' body'

collect :: SCDefn -> Collector ()
collect defn = tell [defn]

We can see that the resulting code is much more readable by removing all the clutters that not not essential to the core logic.

Many applications of monads actually do not require monads but only applicative functors. Monads allow us to run actions depending on the results of earlier actions, but not all applications of monads need this extended functionality.

It is better to use applicative functors where possible because there are some advantages of applicative functors. Applicative functor on the Haskell Wiki mentions two:

• Code that uses only on the Applicative interface are more general than ones uses the Monad interface, because there are more applicative functors than monads. The ZipList is an applicative functor on lists, where liftA2 is implemented by zipWith. It is a typical example of an applicative functor that is not a monad.

• Programming with Applicative has a more applicative/functional feel. Especially for newbies, it may encourage functional style even when programming with effects. Monad programming with do notation encourages a more sequential & imperative style.

There is another advantage. Applicative functors do not need special transformers because they can be combined in a generic way.

But there is a problem. It is usually not easy to decide if we need monads or applicative functors up front. You ambitiously start with applicative functors and find later that you actually needed monads. Sad!

Here is my tip. I start with monads but use only return, ap, liftM, liftM2, … instead of do, >>=. The most common pattern is

do x <- fx
   y <- fy
   return (g x y)

This can be rewritten as liftM2 g fx fy. Once you are sure that you need only those monad methods, you can mechanically switch from monads to applicative functors using the following translation table:

• import Control.Monad -> import Control.Applicative
• return -> pure
• ap -> -> (<*>)
• liftM -> liftA or (<$>)
• liftM2 -> liftA2
• (Monad m =>) -> (Applicative f =>)