A-Normalization example with if-conditionals in OCaml, based on https://matt.might.net/articles/a-normalization/ + alpha conversion

alpha.ml

module A = Ast

let n = ref (-1)
let genid s = incr n; s ^ "." ^ string_of_int !n

let find = List.assoc

let rec g env e = match e with
| A.Int _ -> e
| A.Bool _ -> e
| A.Add(e1,e2) -> A.Add(g env e1, g env e2)
| A.Lt(e1,e2) -> A.Lt(g env e1, g env e2)
| A.Var s -> A.Var(find s env)
| A.Let(s,e1,e2) ->
  let s' = genid s in
  A.Let(s', g env e1, g ((s,s')::env) e2)
| A.If(e1,e2,e3) -> A.If(g env e1, g env e2, g env e3)

let f = g []

anormal.ml

module A = Ast

let n = ref (-1)

let gensym () = incr n; "g" ^ string_of_int !n

let is_value = function
| A.Int _ | A.Bool _ | A.Var _ -> true
| A.Add _ | A.Lt _ | A.Let _ | A.If _ -> false

let id x = x

let rec normalize m k = match m with
| A.Let(v,e1,e2) ->
  normalize e1 (fun x ->
    A.Let(v,x, normalize e2 k))
| A.If(e1,e2,e3) ->
  normalize_name e1 (fun x ->
    k (A.If(x, normalize e2 id, normalize e3 id)))
| A.Add(e1,e2) ->
  normalize_name e1 (fun x ->
    normalize_name e2 (fun y ->
      k (A.Add(x,y))))
| A.Lt(e1,e2) ->
  normalize_name e1 (fun x ->
    normalize_name e2 (fun y ->
      k (A.Lt(x,y))))
| _ -> k m
and normalize_name m k =
  normalize m (fun n ->
    if is_value(n) then k n
    else let g = gensym () in A.Let(g,n,k (A.Var g)))

let f m = normalize m id

ast.ml

type t =
| Int of int
| Bool of bool
| Add of t * t
| Lt of t * t
| Var of string
| Let of string * t * t
| If of t * t * t

let rec to_string = function
| Int n -> string_of_int n
| Bool b -> string_of_bool b
| Add(n, m) -> "(+ " ^ to_string n ^ " " ^ to_string m ^ ")"
| Lt(n, m) -> "(< " ^ to_string n ^ " " ^ to_string m ^ ")"
| Var s -> s
| Let(v,e1,e2) -> "(let [" ^ v ^ " " ^ to_string e1 ^ "] " ^ to_string e2 ^ ")"
| If(e1,e2,e3) -> "(if " ^ to_string e1 ^ " " ^ to_string e2 ^ " " ^ to_string e3 ^ ")"

lexer.mll

{
open Parser
}

let digit = ['0'-'9']
let number = '-'? digit digit*
let whitespace = ['\t' ' ' '\n']

let alpha = ['a'-'z''A'-'Z']
let var = alpha (alpha|digit)*

rule f = parse
  | whitespace+ { f lexbuf }
  | "(" { LPAREN }
  | ")" { RPAREN }
  | "[" { LBRACK }
  | "]" { RBRACK }
  | "+" { PLUS }
  | "<" { LT }
  | "let" { LET }
  | "if" { IF }
  | "true" { TRUE }
  | "false" { FALSE }
  | number as n { INT (int_of_string n ) }
  | var as s { VAR s }
  | eof { EOF }

main.ml

let f s =
  Lexing.from_string s
  |> Parser.f Lexer.f
  |> Alpha.f
  |> Anormal.f
  |> Ast.to_string
  |> print_endline

let () = read_line () |> f

parser.mly

%{
open Ast
%}

%token LPAREN
%token RPAREN
%token LBRACK
%token RBRACK
%token <string> VAR
%token EOF
%token <int> INT
%token PLUS
%token LET
%token TRUE
%token FALSE
%token LT
%token IF

%start <Ast.t> f

%%

f:
  | expr EOF { $1 }

expr:
  | INT { Int $1 }
  | TRUE { Bool true }
  | FALSE { Bool false }
  | VAR { Var $1 }
  | LPAREN PLUS expr expr RPAREN { Add ($3, $4) }
  | LPAREN LT expr expr RPAREN { Lt ($3, $4) }
  | LPAREN LET LBRACK VAR expr RBRACK expr RPAREN { Let($4, $5, $7)}
  | LPAREN IF expr expr expr RPAREN { If($3, $4, $5)}