---
layout: post
title: "Normal Forms"
date: 2016-09-05
comments: true
external-url:
author: Ranjit Jhala
published: true
categories: measures
demo: ANF.hs
---

I have been preparing an undergraduate course on
Compilers in which we build a compiler that crunches
an ML-like language to X86 assembly.
One of my favorite steps in the compilation is the
[conversion to A-Normal Form (ANF)][anf-felleisen]
where, informally speaking, each call or primitive
operation's arguments are **immediate** values,
i.e. constants or variable lookups whose values can
be loaded with a single machine instruction. For example,
the expression

```haskell
((2 + 3) * (12 - 4)) * (7 + 8)
```

has the A-Normal Form:

```haskell
"let anf0 = 2 + 3
     anf1 = 12 - 4
     anf2 = anf0 * anf1
     anf3 = 7 + 8
  in
    anf2 * anf3
```

The usual presentation of ANF conversion
is very elegant but relies upon a clever
use of [continuations][anf-might].
Lets look at another perhaps simpler approach,
where we can use refinements to light the way.

<!-- more -->

<div class="hidden">
\begin{code}
{-@ LIQUID "--no-termination" @-}
{-@ LIQUID "--total"          @-}

module ANF (Op (..), Expr (..), isImm, isAnf, toAnf) where

import Control.Monad.State.Lazy

mkLet :: [(Var, AnfExpr)] -> AnfExpr -> AnfExpr
imm, immExpr :: Expr -> AnfM ([(Var, AnfExpr)], ImmExpr)
anf   :: Expr -> AnfM AnfExpr

-- data IExpr
  -- = IInt Int
  -- | IVar Var
--
-- data AExpr
  -- = AImm IExpr
  -- | ABin Op    IExpr IExpr
  -- | ALet Var   AExpr AExpr
  -- | ALam Var   AExpr
  -- | AApp IExpr IExpr

type AnfExpr = Expr
type ImmExpr = Expr
\end{code}
</div>

Source Language
---------------

Lets commence by defining the source language that we wish to work with:

\begin{code}
data Expr
  = EInt  Int               -- ^ Integers
  | EVar  Var               -- ^ Variables
  | EBin  Op   Expr Expr    -- ^ Binary Operators
  | ELet  Var  Expr Expr    -- ^ Let-binders
  | ELam  Var  Expr         -- ^ Function definitions
  | EApp  Expr Expr         -- ^ Function applications
  deriving (Show)
\end{code}

The language, defined by `Expr` has integers, variables, binary operators,
let-binders and function definitions (lambdas) and calls (applications).
In the above, `Var` and `Op` are simply:

\begin{code}
type Var = String

data Op  = Plus | Minus | Times
         deriving (Show)
\end{code}

For example, the source expression above corresponds to the value:

\begin{code}
-- ((2 + 3) * (12 - 4)) * (7 + 8)
srcExpr :: Expr
srcExpr = EBin Times
            (EBin Times
              (EBin Plus  (EInt  2) (EInt 3))
              (EBin Minus (EInt 12) (EInt 4)))
            (EBin Plus (EInt 7) (EInt 8))
\end{code}

A-Normal Form
-------------

Before we can describe a _conversion_ to A-Normal Form (ANF),
we must pin down what ANF _is_. Our informal description was:
``each call or primitive operation's arguments are immediate
  values, i.e. constants or variable lookups''.

We _could_ define a brand new datatypes, say `IExpr` and `AExpr`
whose values correspond to _immediate_ and _ANF_ terms.
(Try it, as an exercise.) Unfortunately, doing so leads to a
bunch of code duplication, e.g. duplicate printers for `Expr`
and `AExpr`. Instead, lets see how we can use refinements to
carve out suitable subsets.

**Immediate Expressions**

An `Expr` is **immediate** if it is a `Number` or a `Var`;
we can formalize this as a Haskell predicate:

\begin{code}
{-@ measure isImm @-}
isImm :: Expr -> Bool
isImm (EInt _) = True
isImm (EVar _) = True
isImm _        = False
\end{code}

and then we can use the predicate to define a refinement type
for _immediate_ expressions:

\begin{code}
{-@ type ImmExpr = {v:Expr | isImm v} @-}
\end{code}

For example, `e1` is immediate but `e2` is not:

\begin{code}
{-@ e1 :: ImmExpr @-}
e1 = EInt 7

{-@ e2 :: ImmExpr @-}
e2 = EBin Plus e1 e1
\end{code}

**ANF Expressions**

Similiarly, an `Expr` is in **ANF** if all arguments for operators
and applications are **immediate**. Once again, we can formalize
this intuition as a Haskell predicate:

\begin{code}
{-@ measure isAnf @-}
isAnf :: Expr -> Bool
isAnf (EInt {})      = True
isAnf (EVar {})      = True
isAnf (EBin _ e1 e2) = isImm e1 && isImm e2  -- args for operators
isAnf (EApp e1 e2)   = isImm e1 && isImm e2  -- must be immediate,
isAnf (ELet _ e1 e2) = isAnf e1 && isAnf e2  -- and sub-expressions
isAnf (ELam _ e)     = isAnf e               -- must be in ANF
\end{code}

and then use the predicate to define the subset of _legal_ ANF expressions:

\begin{code}
{-@ type AnfExpr = {v:Expr | isAnf v} @-}
\end{code}

For example, `e2` above _is_ in ANF but `e3` is not:

\begin{code}
{-@ e2' :: AnfExpr @-}
e2' = EBin Plus e1 e1

{-@ e3 :: AnfExpr @-}
e3 = EBin Plus e2' e2'
\end{code}

ANF Conversion: Intuition
-------------------------

Now that we have clearly demarcated the territories belonging to plain `Expr`,
immediate `ImmExpr` and A-Normal `AnfExpr`, lets see how we can convert the
former to the latter.

Our goal is to convert expressions like

```haskell
((2 + 3) * (12 - 4)) * (7 + 8)
```

into

```haskell
let anf0 = 2 + 3
    anf1 = 12 - 4
    anf2 = anf0 * anf1
    anf3 = 7 + 8
in
    anf2 * anf3
```

Generalising a bit, we want to convert

```haskell
e1 + e2
```

into

```
let x1  = a1  ... xn  = an
    x1' = a1' ... xm' = am'
in
    v1 + v2
```

where, `v1` and `v2` are immediate, and `ai` are ANF.

**Making Arguments Immediate**

In other words, the key requirement is a way to crunch
arbitrary _argument expressions_ like `e1` into **a pair**

```haskell
([(x1, a1) ... (xn, an)], v1)
```

where

1. `a1...an` are `AnfExpr`, and
2. `v1` is an immediate `ImmExpr`

such that `e1` is _equivalent_ to `let x1 = a1 ... xn = an in v1`.
Thus, we need a function

```haskell
imm :: Expr -> ([(Var, AnfExpr)], ImmExpr)
```

which we will use to **make arguments immediate** thereby yielding
a top level conversion function

```haskell
anf :: Expr -> AnfExpr
```

As we need to generate "temporary" intermediate
binders, it will be convenient to work within a
monad that generates `fresh` variables:

\begin{code}
type AnfM a = State Int a

fresh :: AnfM Var
fresh = do
  n <- get
  put (n+1)
  return ("anf" ++ show n)
\end{code}

Thus, the conversion functions will have the types:

```haskell
anf :: Expr -> AnfM AnfExpr
imm :: Expr -> AnfM ([(Var, AnfExpr)], ImmExpr)
```

ANF Conversion: Code
--------------------

We can now define the top-level conversion thus:

\begin{code}
--------------------------------------------------------------------------------
{-@ anf :: Expr -> AnfM AnfExpr @-}
--------------------------------------------------------------------------------
anf (EInt n) =
  return (EInt n)

anf (EVar x) =
  return (EVar x)

anf (ELet x e1 e2) = do
  a1 <- anf e1
  a2 <- anf e2
  return (ELet x a1 a2)

anf (EBin o e1 e2) = do
  (b1s, v1) <- imm e1
  (b2s, v2) <- imm e2
  return (mkLet (b1s ++ b2s) (EBin o v1 v2))

anf (ELam x e) = do
  a <- anf e
  return (ELam x a)

anf (EApp e1 e2) = do
  (b1s, v1) <- imm e1
  (b2s, v2) <- imm e2
  return (mkLet (b1s ++ b2s) (EApp v1 v2))
\end{code}

In `anf` the real work happens inside `imm` which takes an arbitary
_argument_ expression and makes it **immediate** by generating temporary
(ANF) bindings. The resulting bindings (and immediate values) are
composed by the helper `mkLet` that takes a list of binders and a body
`AnfExpr` and stitches them into a single `AnfExpr`:

\begin{code}
{-@ mkLet :: [(Var, AnfExpr)] -> AnfExpr -> AnfExpr @-}
mkLet []         e' = e'
mkLet ((x,e):bs) e' = ELet x e (mkLet bs e')
\end{code}

The arguments are made immediate by `imm` defined as:

\begin{code}
--------------------------------------------------------------------------------
{-@ imm :: Expr -> AnfM ([(Var, AnfExpr)], ImmExpr) @-}
--------------------------------------------------------------------------------
imm (EInt n)       = return ([], EInt n)
imm (EVar x)       = return ([], EVar x)
imm e@(ELet {})    = immExpr e
imm e@(ELam {})    = immExpr e
imm (EBin o e1 e2) = imm2 e1 e2 (EBin o)
imm (EApp e1 e2)   = imm2 e1 e2 EApp
\end{code}

* Numbers and variables are already immediate, and are returned directly.

* Let-binders and lambdas are simply converted to ANF, and assigned to a fresh
  binder:

\begin{code}
{-@ immExpr :: Expr -> AnfM ([(Var, AnfExpr)], ImmExpr) @-}
immExpr e = do
  a <- anf e
  t <- fresh
  return ([(t, a)], EVar t)
\end{code}

* Finally, binary operators and applications are converted by `imm2` that
  takes two arbitrary expressions and an expression constructor, yielding
  the anf-binders and immediate expression.

\begin{code}
imm2 :: Expr -> Expr -> (ImmExpr -> ImmExpr -> AnfExpr)
     -> AnfM ([(Var, AnfExpr)], ImmExpr)
imm2 e1 e2 f = do
  (b1s, v1) <- imm e1
  (b2s, v2) <- imm e2
  t         <- fresh
  let bs'    = b1s ++ b2s ++ [(t, f v1 v2)]
  return      (bs', EVar t)
\end{code}


You can run it thus:

\begin{code}
toAnf :: Expr -> AnfExpr
toAnf e = evalState (anf e) 0
\end{code}

```haskell
ghci> toAnf srcExpr
ELet "anf0" (EBin Plus (EInt 2) (EInt 3))
 (ELet "anf1" (EBin Minus (EInt 12) (EInt 4))
   (ELet "anf2" (EBin Times (EVar "anf0") (EVar "anf1"))
     (ELet "anf3" (EBin Plus (EInt 7) (EInt 8))
       (EBin Times (EVar "anf2") (EVar "anf3")))))
```

which, can be pretty-printed to yield exactly the outcome desired at the top:

```haskell
let anf0 = 2 + 3
    anf1 = 12 - 4
    anf2 = anf0 * anf1
    anf3 = 7 + 8
in
    anf2 * anf3
```


There! The refinements make this tricky conversion quite
straightforward and natural, without requiring us to
duplicate types and code. As an exercise, can you use
refinements to:

1. Port the classic [continuation-based conversion ?][anf-might]
2. Check that the conversion yields [well-scoped terms ?][lh-scoped]

[lh-scoped]: http://ucsd-progsys.github.io/liquidhaskell-tutorial/10-case-study-associative-maps.html#/using-maps-well-scoped-expressions
[anf-felleisen]: https://users.soe.ucsc.edu/~cormac/papers/pldi93.pdf
[anf-might]: http://matt.might.net/articles/a-normalization/