CEK machine on A-Normalized IR in OCaml, based on https://matt.might.net/articles/cesk-machines/

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)
| A.Fn(ss,e) ->
  let ss' = List.map genid ss in
  A.Fn(ss', g ((List.combine ss ss') @ env) e)
| A.Call(e,es) -> A.Call(g env e, List.map (g env) es)

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 _ | A.Fn _ -> true
| A.Add _ | A.Lt _ | A.Let _ | A.If _ | A.Call _ -> 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))))
| A.Fn(vs,e) ->
  k (A.Fn(vs,normalize e id))
| A.Call(e,es) ->
  normalize_name e (fun x ->
    normalize_names es (fun ys ->
      k (A.Call(x,ys))))
| _ -> 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)))
and normalize_names ms k = match ms with
| [] -> k []
| m::ms ->
  normalize_name m (fun n ->
    normalize_names ms (fun ns -> k (n::ns)))

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
| Fn of string list * t
| Call of t * t list

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 ^ ")"
| Fn(vs,e) ->
  let vs' = String.concat " " vs in
  "(fn [" ^ vs' ^ "] " ^ to_string e ^ ")"
| Call(e,es) ->
  let es' = String.concat " " @@ List.map to_string es in
  "(" ^ to_string e ^ " " ^ es' ^ ")"

cek.ml

module A = Ast

type val_ =
| Int of int
| Bool of bool
| Closure of Ast.t * env
and env = (string * val_) list

type kont =
| Letkont of string * env * Ast.t * kont
| Halt

type cek = 
| Running of Ast.t * env * kont
| Done of val_

let eval_atomic e env = match e with
| A.Int n -> Int n
| A.Bool b -> Bool b
| A.Fn _ -> Closure(e,env)
| A.Var v -> List.assoc v env
| _ -> failwith "Evaluating non-atomic expression as atomic"

let apply_kont k v = match k with
| Halt -> Done v 
| Letkont(s,e,c,k) -> Running(c,(s,v)::e,k)

let step c e k = match c with
| A.Int _ | A.Bool _ | A.Fn _ | A.Var _ -> apply_kont k (eval_atomic c e)
| A.Add(e1,e2) -> (match eval_atomic e1 e, eval_atomic e2 e with
  | Int n, Int m -> apply_kont k (Int(n+m))
  | _ -> failwith "Adding non-integer values")
| A.Lt(e1,e2) -> (match eval_atomic e1 e, eval_atomic e2 e with
  | Int n, Int m -> apply_kont k (Bool(n<m))
  | _ -> failwith "Comparing non-integer values")
| A.If(cond,e1,e2) -> (match eval_atomic cond e with
  | Bool b -> Running((if b then e1 else e2), e, k)
  | _ -> failwith "Conditional on non-boolean")
| A.Let(s,e1,e2) -> Running(e1,e,Letkont(s,e,e2,k))
| A.Call(f,es) -> (match eval_atomic f e with
  | Closure(A.Fn(ss,body),e') ->
    let vs = List.map (fun v -> eval_atomic v e) es in
    Running(body,(List.combine ss vs)@e', k)
  | _ -> failwith "Non-function in operator position")

let eval c =
  let rec f = function
  | Running(c,e,k) -> f(step c e k)
  | Done v -> v
  in
  f (Running(c,[],Halt))

let string_of_val = function
| Int n -> string_of_int n
| Bool b -> string_of_bool b
| Closure _ -> "closure"

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 }
  | "fn" { FN }
  | "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

let () =
  let e = read_line() |> f in
  e |> L.Ast.to_string |> print_endline;
  e |> L.Cek.eval |> L.Cek.string_of_val |> print_endline;

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
%token FN

%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)}
  | LPAREN FN LBRACK list(VAR) RBRACK expr RPAREN { Fn($4, $6)}
  | LPAREN expr list(expr) RPAREN { Call($2, $3) }