A-normalization (administrative normal form, ANF)

ANF.hs

module ANF where

-- https://www.researchgate.net/publication/2392886_The_Essence_of_Compiling_with_Continuations
-- https://matt.might.net/articles/a-normalization/
-- https://gist.github.com/siraben/447c419508b460afc99a232d5df8063f

import Control.Monad.State

type Name = String

data SExpr
  = SVar Name
  | SLam Name SExpr
  | SApp SExpr SExpr
  | SLet Name SExpr SExpr
  deriving (Eq, Show)

-- Administrative normal form (ANF) expressions have to satisfy the constraints
-- imposed by the two grammars below. Each grammar reveals some structure not
-- so apparent in the other.
--
-- Grammar based on The Essence of Compiling with Continuations:
--
--  <expr> ::= <value>
--           | 'let' <var> '=' <value> 'in' <expr>
--           | <value> <value>
--           | 'let' <var> '=' <value> <value> 'in' <expr>
--
--  <value> ::= <var>
--            | 'λ' <var> ',' <expr>
--
-- Let expressions can only bind variables, applications and primitive
-- operations (not present in the grammar above), in particular, let
-- expressions cannot be let bound.
--
-- Grammar based on Matthew Might's blog post:
--
--  <aexpr> ::= <var>
--            | 'λ' <var> ',' <expr>
--
--  <cexpr> ::= <aexpr> <aexpr>
--
--  <expr> ::= 'let' <var> '=' <expr> 'in' <expr>
--           | <aexpr>
--           | <cexpr>
--
-- Evaluating an atomic expression (`<aexpr>`) is guaranteed to terminate,
-- cause no side effects, cause no control effects and never produce an error.
-- A lambda expression is atomic because evaluating it creates a closure and
-- that process satisfies the required constraints.
--
-- Evaluating a complex expression (`<cexpr>`) may not terminate, may have a
-- side effect and may produce an error. However, a complex expression may
-- defer execution to only one other complex expression. For example, the
-- following are also complex expressions:
--
--  <cexpr> ::= <aexpr> <aexpr>
--            | 'if' <aexpr> 'then' <expr> 'else' <expr>
--            | 'let' 'rec' <var> '=' <aexpr> 'in' <expr>
--
-- Although function application can apply any two values, the normalizer below
-- only allows applying variables to variables. This is similar to the
-- continuation passing style (CPS) grammars in Compiling with Continuations,
-- Continued.

data AExpr
  = Var Name
  | Lam Name Expr
  deriving (Eq, Show)

data CExpr
  = App Name Name
  deriving (Eq, Show)

data Expr
  = Let Name Expr Expr
  | AExpr AExpr
  | CExpr CExpr
  deriving (Eq, Show)

fresh :: State Int String
fresh = state $ \x -> (show x, x + 1)

normalize :: SExpr -> Expr
normalize e = evalState (expr e) 0
  where
  expr :: SExpr -> State Int Expr
  expr (SVar x) = return $ AExpr (Var x)
  expr (SLam x b) = AExpr . Lam x <$> expr b
  expr (SApp f a) =
    name f $ \f ->
      name a $ \a ->
        return $ CExpr (App f a)
  expr (SLet x e1 e2) = do
    e1 <- expr e1
    e2 <- expr e2
    -- As pointed out in Compiling with Continuations, Continued, the
    -- A-normalization algorithm in The Essence of Compiling with Continuations
    -- doesn't actually normalizes terms, as it doesn't deal with the let
    -- expression invariant.
    case e1 of
      Let x' e1' e2' -> return $ Let x' e1' (Let x e2' e2)
      _ -> return $ Let x e1 e2

  name :: SExpr -> (Name -> State Int Expr) -> State Int Expr
  name e k = do
    e <- expr e
    case e of
      AExpr (Var x) -> k x
      _ -> do
        x <- fresh
        Let x e <$> k x