ANF transformation with joint point optimization for branching, using explicit CPS

ANF-with-joint-point.ml

module Src = struct
    type expr =
        | Int of int
        | Var of string
        | Op  of string * expr * expr
        | Let of string * expr * expr
        | If  of expr * expr * expr
end


module IR = struct
    type variable = int
    let var_seed = ref 0
    let fresh_var () = incr var_seed; !var_seed

    type label = int
    let label_seed = ref 0
    let fresh_label () = incr label_seed; !label_seed

    type value =
        | Int of int
        | Var of variable

    type expr =
        | Val of value
        | Op  of string * value * value

    type program =
        | Ret of value
        | Jmp of label * value
        | Blk of (label * variable * program) * program
        | Let of variable * expr * program
        | If  of value * program * program
end


type continuation =
    | Ret
    | Blk of IR.label
    | Fun of (IR.value -> IR.program)


let apply (type a) k v =
    match k with
    | Ret     -> IR.Ret v
    | Blk lbl -> IR.Jmp(lbl, v)
    | Fun k   -> k v


module Env = Map.Make(String)

let (let*) f k = f (Fun k)

let rec tyck env expr k =
    match expr with
    | Src.Int i -> apply k (IR.Int i)
    | Src.Var v -> apply k (Env.find v env)
    | Src.Op(op, l, r) ->
        let* lv = tyck env l in
        let* rv = tyck env r in
        let var = IR.fresh_var () in
        IR.Let(var, IR.Op(op, lv, rv), apply k (IR.Var var))
    | Src.Let(name, rhs, body) ->
        let* rhsv = tyck env rhs in
        let var = IR.fresh_var () in
        IR.Let(var, IR.Val rhsv, apply k (IR.Var var))
    | Src.If(cond, conseq, alter) ->
        let* condv = tyck env cond in
        let lbl = IR.fresh_label () in
        let param = IR.fresh_var () in
        IR.Blk( (lbl, param, apply k (IR.Var param))
             , IR.If( condv
                   , tyck env conseq (Blk lbl)
                   , tyck env alter  (Blk lbl) ))

let test expr = tyck Env.empty expr Ret


let test1 =
    let open Src in
    Op("add", Int 1, Op("mul", Int 2, Int 3))

let test2 =
    let open Src in
    Let( "x", Op("add", Int 1, Int 2)
       , Let( "y", Op("mul", Int 3, Int 4)
            , Op("mul", Var "x", Op("add", Var "y", Var "y") )))

let test3 =
    let open Src in
    Op("add", Int 1
      , If( Int 0
          , Op("mul", Int 2, Int 3)
          , Op("mul", Int 4, Int 5) ))

let test4 =
    let open Src in
    If( Int 0
      , Op("mul", Int 2, Int 3)
      , Op("mul", Int 4, Int 5) )