Monadic Normal Form transform using shift/reset.

mnf.ml

open Delimcc;;

type id = string

type exp =
  | Var of id
  | Fun of id * exp
  | Const of int
  | BinOp of exp * exp
  | App of exp * exp

type value =
  | VVar of id
  | VFun of id * mnf
  | VConst of int

and r_exp =
  | RBinOp of value * value
  | RApp of value * value

and mnf =
  | Ret of value
  | Let of id * r_exp * mnf

let counter = ref 0

let new_id prefix : id =
  let s = prefix ^ (string_of_int !counter) in
  counter := !counter + 1;
  s

let new_var () = new_id "var:"

let p = new_prompt ()

let rec conv_exp e : mnf =
  push_prompt p (fun () -> Ret (conv_exp' e))
and conv_exp' e : value =
  match e with
    | BinOp (e1, e2) ->
      shift p (fun k ->
	let v1 = conv_exp' e1 in
	let v2 = conv_exp' e2 in 
	let x = new_var () in
	Let (x, RBinOp (v1, v2), k (VVar x)))
    | App (e1, e2) -> 
      shift p (fun k ->
	let v1 = conv_exp' e1 in
	let v2 = conv_exp' e2 in 
	let x = new_var () in
	Let (x, RApp (v1, v2), k (VVar x)))
    | Var v -> VVar v
    | Fun (x, e) -> VFun (x, conv_exp e)
    | Const i -> VConst i

(*
  \x. ((x + 3) + ((\y. 2 + y) x)) + 1
*)
let p = Fun ("x",
	     BinOp (BinOp (BinOp (Var "x", Const 3),
			   App (Fun ("y", BinOp (Const 2, Var "y")),
				Var "x")),
		    Const 1))

let rec conv_exp_cps e : mnf =
  conv_exp_cps' e (fun x -> Ret x)
and conv_exp_cps' e k =
  match e with
    | BinOp (e1, e2) ->
      conv_exp_cps' e1 (fun v1 ->
      conv_exp_cps' e2 (fun v2 ->
	let x = new_var () in
	Let (x, RBinOp (v1, v2), k (VVar x))))
    | App (e1, e2) ->
      conv_exp_cps' e1 (fun v1 ->
      conv_exp_cps' e2 (fun v2 ->
	let x = new_var () in
	Let (x, RApp (v1, v2), k (VVar x))))
    | Var v -> k (VVar v)
    | Fun (x, e) -> k (VFun (x, conv_exp_cps e))
    | Const i -> k (VConst i)

let rec string_of_mnf = function
  | Ret v -> "Ret " ^ string_of_value v
  | Let (x, e1, e2) -> "let " ^ x ^ " = " ^ string_of_r_exp e1 ^ " in\n" ^ string_of_mnf e2
and string_of_value : value -> string = function
  | VVar x -> x
  | VFun (x, e) -> 
    "(fun " ^ x ^ " ->\n" ^ string_of_mnf e ^ ")"
  | VConst i -> string_of_int i
and string_of_r_exp = function
  | RBinOp (v1, v2) -> string_of_value v1 ^ " + " ^ string_of_value v2
  | RApp (v1, v2) -> string_of_value v1 ^ " " ^ string_of_value v2

let print_mnf e = print_endline (string_of_mnf e)

(*
Ret (fun x ->
  let a = x + 3 in
  let b = (fun y -> let c = 2 + y in c) x in
  let d = a + b in
  let e = d + 1 in
  Ret e)
*)
let _ = print_mnf (conv_exp p)
let _ = print_mnf (conv_exp_cps p)