Implementing Functional Languages: a tutorial, Template instantiation

compiler.ml

open Syntax

module IntMap = Map.Make(struct type t = int let compare = compare end)

module List = struct
  include List
  let zip xs ys =
    let rec go acc = function
      | [], _ | _, [] -> List.rev acc
      | x::xs, y::ys -> go ((x,y)::acc) (xs,ys)
    in
    go [] (xs,ys)
  let rec take n xs =
    match n, xs with
    | 0, _ | _, [] -> []
    | n, x::xs -> x :: take (n-1) xs
  let rec drop n xs =
    match n, xs with
    | 0, _ | _, [] -> xs
    | n, _::xs -> drop (n-1) xs
  let rec last = function
    | [] -> failwith "last: empty list"
    | [x] -> x
    | x::xs -> last xs
  let split_at k xs =
    let rec go l k = function
      | _ when k = 0 ->
          l, []
      | [] ->
          l, []
      | x::xs ->
          go (x::l) (k-1) xs
    in
    go [] k xs
end

type addr = int
type mark = int
type primitive =
  | Neg | Add | Sub | Mul | Div
  | Lt | Le | Gt | Ge | Eq | Ne
  | PrimConstr of int * int
  | Stop
  | Print
type node =
  | NAp of addr * addr
  | NSc of name * name list * core_expr
  | NNum of int
  | NInd of addr
  | NPrim of primitive
  | NData of int * addr list
  | NMark of mark * node

type ti_stack = int list
type ti_dump = ti_stack list
type ti_heap = node IntMap.t
type ti_globals = (name * addr) list
type ti_stats = int
type tiState = ti_stack * ti_dump * ti_heap * ti_globals * ti_stats

let prelude = Parser.program Lexer.token (Lexing.from_string
"
I x = x;
K x y = x;
K1 x y = y;
S f g x = f x (g x);
compose f g x = f (g x);
twice f = compose f f;

False t f = f;
True t f = t;
if = I;
and b1 b2 t f = b1 (b2 t f) f;
or b1 b2 t f = b1 t (b2 t f);
not b t f = b f t;

pair a b f = f a b;
fst p = p K;
snd p = p K1;

cons a b cn cc = cc a b;
nil cn cc = cn;
empty l = l True (K (K False));
head l = l stop K;
tail l = l stop K1;

printList xs = xs stop printCons;
printCons h t = print h (printList t);
")

let map_accuml_rev f acc xs =
  let acc, ys = List.fold_left (fun (acc,ys) x ->
    let acc, y = f acc x in
    acc, y::ys
  ) (acc,[]) xs
  in
  acc, ys

let map_accuml f acc xs =
  let acc, ys = map_accuml_rev f acc xs in
  acc, List.rev ys

let heap_counter = ref 0

let heap_alloc heap obj =
  let id = !heap_counter in
  incr heap_counter;
  IntMap.add id obj heap, id

let heap_update heap addr obj = IntMap.add addr obj heap

let heap_update_alloc heap addr obj =
  if addr < 0 then
    heap_alloc heap obj
  else
    heap_update heap addr obj, addr

let heap_lookup heap addr =
  try
    IntMap.find addr heap
  with Not_found ->
    Printf.eprintf "heap_lookup: %d not found\n" addr;
    raise Not_found

let heap_free heap addr = IntMap.remove addr heap

let a_lookup assoc key =
  let rec go = function
    | [] -> failwith (Printf.sprintf "a_lookup: key %s not found" key)
    | (k,v)::xs when k = key -> v
    | _::xs -> go xs
  in
  go assoc

let alloc_sc heap (name, args, body) =
  let heap, addr = heap_alloc heap (NSc (name, args, body)) in
  heap, (name, addr)

let primitives = [
  "negate", Neg;
  "print", Print;
  "stop", Stop;
  "+", Add;
  "-", Sub;
  "*", Mul;
  "/", Div;
  "==", Eq;
  "!=", Ne;
  "<", Lt;
  "<=", Le;
  ">", Gt;
  ">=", Ge
]

let alloc_prim heap (name, prim) =
  let heap, addr = heap_alloc heap (NPrim (a_lookup primitives name)) in
  heap, (name, addr)

let compile program =
  let heap, sc_addrs = map_accuml_rev alloc_sc IntMap.empty (prelude @ program) in
  let heap, prim_addrs = map_accuml_rev alloc_prim heap primitives in
  let init_stat = 0 in
  [a_lookup sc_addrs "main"], [], heap, sc_addrs @ prim_addrs, init_stat

(* show *)

open Pprint

let column =
  Column (fun k -> if k < 60 then Text (String.make (60-k) ' ') else Line true)

let rec show_node heap = function
  | NNum n -> fill 5 (Text "NNum") <+> int n
  | NSc (name,args,body) -> fill 5 (Text "NSc") <+> Text name
  | NAp (a1,a2) -> fill 5 (Text "NAp") <+> int a1 <+> int a2
  | NInd a -> fill 5 (Text "NInd") <+> int a
  | NPrim prim -> Text "NPrim" <+>
    (match prim with
    | Neg -> Text "Neg"
    | Add -> Text "Add"
    | Sub -> Text "Sub"
    | Mul -> Text "Mul"
    | Div -> Text "Div"
    | Eq -> Text "Eq"
    | Ne -> Text "Ne"
    | Lt -> Text "Lt"
    | Le -> Text "Le"
    | Gt -> Text "Gt"
    | Ge -> Text "Ge"
    | PrimConstr (tag,n) -> Text "Constr" <+> int tag <+> int n
    | Stop -> Text "Stop"
    | Print -> Text "Print"
    )
  | NData (tag,args) ->
      (match tag with
      | _ ->
          Text "NData" <+> int tag <+> sep_by (Char ',') (List.map int args))
  | NMark (m,o) ->
      Text "NMark" <+> int m <+> show_node heap o

let show_stack heap stack =
  let show_stk_node = function
    | NAp (a1, a2) ->
        fill 4 (Text "NAp") <+> int a1 <+> int a2 <+>
        enclose (Char '(') (Char ')') (show_node heap @@ heap_lookup heap a2)
    | node -> show_node heap node
  in
  let show_stack_item addr =
    intw 2 addr <.> Char ':' <+> show_stk_node (heap_lookup heap addr)
  in
  Text "s" <+> Char '[' <.> sep_by (Line true)
  (List.map show_stack_item stack)
  <.> Char ']'

let show_heap heap =
  Text "heap" <+>
  (*Char '[' <.> sep_by (Line true)*)
  Char '[' <.> sep_by column
  (List.map (fun (i,node) -> intw 2 i <.> Char ')' <+> show_node heap node) @@ IntMap.bindings heap)
  <.> Char ']'

let show_dump dump =
  Text "dump" <+>
  enclose_sep (Char '[') (Char ']') (Line false <.> Text ", ") @@
  List.map (fun d -> enclose_sep (Char '[') (Char ']') (Char ',') @@ List.map int d) dump

let show_globals globals =
  Text "globals" <+>
  enclose_sep (Char '[') (Char ']') (Text ", ") @@
  List.map (fun (name,addr) -> Text name <+> int addr) globals

let show_state (stack, dump, heap, globals, stats) =
  show_stack heap stack <.> Column (fun k -> if k < 36 then Text (String.make
  (36-k) ' ') else Empty) <.> show_heap heap
  <$>
  show_dump dump
  <$>
  show_globals globals

let show_stats (stack, dump, heap, globals, stats) =
  Line false <.> Text "total number of steps =" <+> int stats

let show_results states =
  Pprint.pp_list (List.map show_state states) <$> show_stats (List.last states)

let backtrace heap node =
  let found = Hashtbl.create 0 in
  let rec go prec x =
    if Hashtbl.mem found x then
      int x
    else (
      Hashtbl.replace found x ();
      match heap_lookup heap x with
      | NNum n -> int n
      | NSc (name,args,body) -> Text name
      | NAp (a1,a2) -> paren (prec > 0) (Text "NAp" <+> go 1 a1 <+> go 1 a2)
      | NInd a -> paren (prec > 0) (Text "NInd" <+> go 1 a)
      | NPrim prim ->
        (match prim with
        | Neg -> Text "Neg"
        | Add -> Text "Add"
        | Sub -> Text "Sub"
        | Mul -> Text "Mul"
        | Div -> Text "Div"
        | Eq -> Text "Eq"
        | Ne -> Text "Ne"
        | Lt -> Text "Lt"
        | Le -> Text "Le"
        | Gt -> Text "Gt"
        | Ge -> Text "Ge"
        | PrimConstr (tag,n) -> paren (prec > 0) (Text "Constr" <+> int tag <+> int n)
        | Stop -> Text "Stop"
        | Print -> Text "Print"
        )
      | NData (tag,args) ->
          paren (prec > 0) (Text "NData" <+> int tag <+> enclose_sep (Char '(')
          (Char ')') (Text ", ")
          (List.map (fun i -> go 0 i) args))
      | NMark (m,o) ->
          paren (prec > 0) (Text "NMark" <+> int m)
    )
  in
  Text "trace" <+> go 0 node

(* eval *)

let is_data_node = function
  | NNum _ -> true
  | NData _ -> true
  | _ -> false

let is_mark_node = function
  | NMark _ -> true
  | _ -> false

let ti_final (stack, dump, heap, globals, stats) =
  dump = [] &&
  (match stack with
  | [] -> true
  | [x] -> is_data_node (heap_lookup heap x)
  | _ -> false)

let get_args heap stack n =
  let rec go acc i xs =
    if i = 0 then
      List.rev acc
    else
      match xs with
      | [] -> assert false
      | x::xs ->
          (match heap_lookup heap x with
          | NAp (a1, a2) -> go (a2::acc) (i-1) xs
          | _ -> assert false)
  in
  go [] n stack

let rec instantiate addr body heap env =
  match body with
  | ENum n ->
      heap_update_alloc heap addr (NNum n)
  | EVar v ->
      let a = a_lookup env v in
      if addr < 0 then
        heap, a
      else
        heap_update_alloc heap addr (NInd a)
  | EConstr (tag,n) ->
      heap_update_alloc heap addr (NPrim (PrimConstr (tag,n)))
  | ELet (false, defs, body) ->
      let heap, binds = List.fold_left (fun (heap,binds) (name,e) ->
          let heap, a = instantiate (-1) e heap env in
          heap, (name,a)::binds
        ) (heap,[]) defs in
      instantiate addr body heap (binds @ env)
  | ELet (true, defs, body) ->
      let heap, binds = List.fold_left (fun (heap,binds) (name,e) ->
          let heap, a = instantiate (-1) (ENum 0) heap env in
          heap, (name,a,e)::binds
        ) (heap,[]) defs in
      let env = List.map (fun (name,a,e) -> name,a) binds @ env in
      let heap = List.fold_left (fun heap (name,a,e) ->
          let heap, a' = instantiate (-1) e heap env in
          let heap = heap_update heap a (heap_lookup heap a') in
          heap_free heap a'
        ) heap binds in
      instantiate addr body heap env
  | ECase _ -> failwith "cannot instantiate case"
  | ELam _ -> failwith "cannot instantiate lambda"
  | EAp (e1,e2) ->
      let heap, a1 = instantiate (-1) e1 heap env in
      let heap, a2 = instantiate (-1) e2 heap env in
      heap_update_alloc heap addr (NAp (a1,a2))

let step_check (stack, dump, heap, globals, stats) n =
  let rec go args root i xs =
    if i = 0 then
      None, (List.rev args, xs, root)
    else
      (match xs with
      | [] ->
          failwith "number of arguments"
      | a::xs' ->
          (match heap_lookup heap a with
          | NAp (_, b) ->
            let obj = heap_lookup heap b in
            if is_data_node obj then
              go (obj::args) a (i-1) xs'
            else
              Some ([b], xs::dump, heap, globals, stats), ([], [], -1)
          | _ ->
              failwith "number of arguments"
          )
      )
  in
  go [] (-1) n (List.tl stack)

let rec step ((stack, dump, heap, globals, stats) as state) =
  let stats = stats + 1 in
  match stack with
  | [] -> failwith "step: empty stack"
  | hd::tl ->
      match heap_lookup heap hd with
      | NNum _ | NData _ ->
          if dump = [] then
            failwith "step: number applied as a function"
          else
            List.hd dump, List.tl dump, heap, globals, stats
      | NAp (a1, a2) ->
          (match heap_lookup heap a2 with
          | NInd a3 ->
              let heap = heap_update heap hd (NAp (a1, a3)) in
              stack, dump, heap, globals, stats
          | _ ->
            a1::stack, dump, heap, globals, stats)
      | NInd a -> a::tl, dump, heap, globals, stats
      | NSc (name, arg_names, body) ->
          let stack' = List.drop (List.length arg_names) stack in
          let heap, addr = instantiate (List.hd stack') body heap
          (List.zip arg_names (get_args heap tl (List.length arg_names)) @ globals) in
          (*let heap = heap_update heap (List.hd stack') (NInd addr) in*)
          (addr::List.tl stack', dump, heap, globals, stats)
      | NPrim Neg ->
          (match step_check state 1 with
          | Some state, _ ->
              state
          | None, ([NNum n], stack, root) ->
              let heap = heap_update heap root (NNum (-n)) in
              root::stack, dump, heap, globals, stats
          | _ ->
              assert false
          )
      | NPrim (Add as prim)
      | NPrim (Sub as prim)
      | NPrim (Mul as prim)
      | NPrim (Div as prim) ->
          (match step_check state 2 with
          | Some state, _ ->
              state
          | None, ([NNum n1; NNum n2], stack, root) ->
              let obj = NNum (match prim with | Add -> n1 + n2 | Sub -> n1 - n2
                                              | Mul -> n1 * n2 | Div -> n1 / n2
                                              | _ -> assert false) in
              let heap = heap_update heap root obj in
              root::stack, dump, heap, globals, stats
          | _ ->
              assert false
          )
      | NPrim (Eq as prim)
      | NPrim (Ne as prim)
      | NPrim (Lt as prim)
      | NPrim (Le as prim)
      | NPrim (Gt as prim)
      | NPrim (Ge as prim) ->
          (match step_check state 2 with
          | Some state, _ ->
              state
          | None, ([NNum n1; NNum n2], stack, root) ->
              let obj = (match prim with | Eq -> n1=n2 | Ne -> n1<>n2 | Lt -> n1<n2
                                         | Le -> n1<=n2 | Gt -> n1>n2 | Ge -> n1>=n2
                                         | _ -> assert false) in
              let heap = heap_update heap root (NInd (a_lookup globals (if obj then "True" else "False"))) in
              root::stack, dump, heap, globals, stats
          | _ ->
              assert false
          )
      | NPrim (PrimConstr (tag,n)) ->
          let xs = List.drop n stack in
          let heap = heap_update heap (List.hd xs) (NData (tag, get_args heap tl n)) in
          xs, dump, heap, globals, stats
      | NPrim Stop ->
          tl, dump, heap, globals, stats
      | NPrim Print ->
          (match tl with
          | a1::a2::xs ->
            (match heap_lookup heap a1 with
            | NAp (_, b) ->
                let r = match heap_lookup heap a2 with
                  | NAp (_, c) ->
                      c::xs, dump, heap, globals, stats
                  | _ ->
                      failwith "step: NPrim Print"
                in
                (match heap_lookup heap b with
                | NNum n ->
                    Printf.printf "%d\n" n;
                    r
                | NData (1,[]) ->
                    Printf.printf "%b\n" false;
                    r
                | NData (2,[]) ->
                    Printf.printf "%b\n" true;
                    r
                | _ ->
                    [b], (a2::xs)::dump, heap, globals, stats
                )
            | _ ->
                failwith "step: NPrim Print"
            )
          | _ ->
              failwith "step: NPrim Print"
          )

(* garbage collection *)

let log = ref true

type ('a,'b) either = Left of 'a | Right of 'b

let mark_from heap addr =
  (*if !log then*)
    (*Printf.eprintf "\n\n-----------\n";*)
  let rec go heap f b =
    let rec marked a =
      match heap_lookup heap a with
      | NMark _ -> Right a
      | NInd o -> marked o
      | _ -> Left a
    in
    let obj = heap_lookup heap f in
    match obj with
    | NMark (-1, _) ->
        if b < 0 then
          heap, f
        else
          (match heap_lookup heap b with
          | NMark (1, NAp (a1,a2)) ->
              (*if !log then*)
                (*Printf.eprintf "+ NAp   %d %3d %3d %d,%d\n" 1 f b a1 a2;*)
              let heap = heap_update heap b (NMark (1, (NAp (f,a2)))) in
              go heap b a1
          | NMark (2, NAp (a1,a2)) ->
              (*if !log then*)
                (*Printf.eprintf "+ NAp   %d %3d %3d %d,%d\n" 2 f b a1 a2;*)
              let heap = heap_update heap b (NMark (2, (NAp (a1,f)))) in
              go heap b a2
          | NMark (mark, NData (tag,a1)) ->
              let revl, r = List.split_at (mark-1) a1 in
              let heap = heap_update heap b (NMark (mark, NData (tag, List.rev (f::revl) @ List.tl r))) in
              go heap b (List.hd r)
          | _ ->
              (*Printf.eprintf "!!! %d %d\n" f b;*)
      (*show_node heap (heap_lookup heap f) |> Pprint.layout 1.0 80 |> prerr_endline;*)
      (*show_node heap (heap_lookup heap b) |> Pprint.layout 1.0 80 |> prerr_endline;*)
              assert false
          )
    | NMark (mark,_) ->
        (match obj with
        | NMark (1, NAp (a1, a2)) ->
            (*if !log then*)
              (*Printf.eprintf "- NAp   %d %3d %3d %d,%d\n" 1 f b a1 a2;*)
            (match marked a2 with
            | Left a2 ->
              let heap = heap_update heap f (NMark (2, NAp (a1, b))) in
              go heap a2 f
            | Right a2 ->
              let heap = heap_update heap f (NMark (2, NAp (a1, a2))) in
              go heap f b
            )
        | NMark (2, NAp (a1, a2)) ->
            (*if !log then*)
              (*Printf.eprintf "- NAp   %d %3d %3d %d,%d\n" 2 f b a1 a2;*)
            let heap = heap_update heap f (NMark (-1, NAp (a1, a2))) in
            go heap f b
        | NMark (mark, (NData (tag,a1) as o)) ->
            let revl, r = List.split_at mark a1 in
            if r = [] then
              let heap = heap_update heap f (NMark (-1, o)) in
              go heap f b
            else
              (match marked (List.hd r) with
              | Left x ->
                let heap = heap_update heap f (NMark (mark+1, NData (tag, List.rev (b::revl) @ List.tl r))) in
                go heap (List.hd r) f
              | Right x ->
                let heap = heap_update heap f (NMark (mark+1, NData (tag, List.rev (x::revl) @ List.tl r))) in
                go heap f b
              )
        | _ ->
            assert false
        )
    | NNum _ | NSc _ | NPrim _ ->
        (*if !log then*)
          (*Printf.eprintf "- N___    %3d %3d\n" f b;*)
        let heap = heap_update heap f (NMark (-1, obj)) in
        go heap f b
    | NInd a ->
        go heap a b
    | NAp (a1,a2) ->
        (*if !log then*)
          (*Printf.eprintf "- NAp     %3d %3d %d,%d\n" f b a1 a2;*)
        (match marked a1 with
        | Left a1 ->
            let heap = heap_update heap f (NMark (1, NAp (b, a2))) in
            go heap a1 f
        | Right a1 ->
            let heap = heap_update heap f (NMark (1, NAp (a1, a2))) in
            go heap f b
        )
    | NData (tag,a1) ->
        (*if !log then*)
          (*Printf.eprintf "- NData   %d %d\n" f b;*)
        (match a1 with
        | [] ->
            let heap = heap_update heap f (NMark (-1, obj)) in
            go heap f b
        | x::xs ->
            (match marked x with
            | Left x ->
              let heap = heap_update heap f (NMark (1, NData (tag, b::xs))) in
              go heap x f
            | Right x ->
              let heap = heap_update heap f (NMark (1, NData (tag, x::xs))) in
              go heap f b
            )
        )
  in
  go heap addr (-1)

let mark_from_stack heap stack =
  map_accuml mark_from heap stack

let mark_from_dump heap dump =
  map_accuml mark_from_stack heap dump

let mark_from_globals heap globals =
  let names, addrs = List.split globals in
  let heap, addrs = map_accuml mark_from heap addrs in
  heap, List.zip names addrs

let scan_heap heap =
  IntMap.fold (fun addr obj heap' ->
    match obj with
    | NMark (m,o) when m < 0 ->
        IntMap.add addr o heap'
    | _ ->
        Printf.eprintf "++ remove %d\n" addr;
        heap'
  ) heap IntMap.empty

(* eval *)

let eval ((stack, dump, heap, globals, stats) as s) =
  let rec go acc ((stack,dump,heap,globals,stats) as s) =
    let s =
      if stats mod 100 = 99 then
        let heap, stack = mark_from_stack heap stack in
        let heap, dump = mark_from_dump heap dump in
        let heap, globals = mark_from_globals heap globals in
        let heap = scan_heap heap in
        stack,dump,heap,globals,stats
      else
        s
    in

    (*Printf.eprintf "::: %d\n" stats;*)
    (*if stats = 29 then 
      (*List.iter (Printf.eprintf "== %d\n") stack;*)
      (*show_node heap (heap_lookup heap 23) |> Pprint.layout 1.0 80 |> prerr_endline;*)
      (*show_node heap (heap_lookup heap 72) |> Pprint.layout 1.0 80 |> prerr_endline;*)
      (*show_node heap (heap_lookup heap 73) |> Pprint.layout 1.0 80 |> prerr_endline;*)
      (*show_node heap (heap_lookup heap 74) |> Pprint.layout 1.0 80 |> prerr_endline;*)
    *)

    intw 3 stats <.> Char ')' <+> show_state s |> Pprint.layout 1.0 80 |> prerr_endline;
    if stack <> [] then
      backtrace heap (List.last stack) |> Pprint.layout 1.0 80 |> prerr_endline;
    if ti_final s then
      List.rev (s::acc)
    else
      go (s::acc) (step s)
  in
  go [] s

lexer.mll

{
open Parser
}
let space = [' ' '\t' '\n' '\r']
let digit = ['0'-'9']
let alpha = ['a'-'z''A'-'Z']

rule token = parse
| space+ { token lexbuf }
| "--" { comment lexbuf; token lexbuf }
| digit+ { INT(Lexing.lexeme lexbuf |> int_of_string) }
| "==" { EQ }
| "!=" { NE }
| '<' { LT }
| "<=" { LE }
| '>' { GT }
| ">=" { GE }
| ';' { SEMICOLON }
| '(' { LPAREN }
| ')' { RPAREN }
| '{' { LBRACKET }
| '}' { RBRACKET }
| ';' { SEMICOLON }
| '*' { MUL }
| '/' { DIV }
| '+' { ADD }
| '-' { SUB }
| '=' { ASSIGN }
| '<' { L }
| '>' { R }
| '(' { LPAREN }
| ')' { RPAREN }
| '\\' { LAMBDA }
| '.' { DOT }
| ',' { COMMA }
| "->" { ARROW }

| "Pack" { PACK }
| "let" { LET }
| "in" { IN }
| "letrec" { LETREC }
| "case" { CASE }
| "of" { OF }

| alpha (digit|alpha|'_'|'\'')* { IDENT(Lexing.lexeme lexbuf) }

| eof { EOF }
| _ { Printf.eprintf "characters %d-%d:\nerror: unknown token %s\n"
      (Lexing.lexeme_start lexbuf)
      (Lexing.lexeme_end lexbuf)
      (Lexing.lexeme lexbuf); exit 1 }
and comment = parse
| '\n' { () }
| _ | eof { comment lexbuf }

main.ml

open Pprint

let width = ref 80
let rfrac = ref 1.0

let channel stage ic =
  try
    (*let show x = Pprint.layout !rfrac !width x |> prerr_string in*)
    let ast = Parser.program Lexer.token (Lexing.from_channel ic) in
    (*Pprint.pp_program ast |> show;*)
    (*print_endline "++++++";*)
    let state = Compiler.compile ast in
    (*Compiler.show_state state |> show*)
    Compiler.eval state |> ignore
    (*Compiler.show_results results |> Pprint.layout !rfrac !width |> print_endline*)
  with e ->
    close_in ic;
    raise e

let file stage f =
  channel stage (if f = "-" then stdin else open_in f)

let () =
  let stage = ref 0 in
  let files = ref [] in
  Arg.parse
    [("-d", Arg.Int(fun i -> stage := i), "stage");
     ("-w", Arg.Int(fun i -> width := i), "width")]
    (fun s -> files := s :: !files)
    ("core compiler");
  if !files = [] then
    channel !stage stdin
  else
    List.rev !files |> List.iter (file !stage)

parser.mly

%{
open Syntax
%}

%token <int> INT
%token <string> IDENT
%token MUL DIV ADD SUB ASSIGN L R LPAREN RPAREN LBRACKET RBRACKET PACK EQ NE LT LE GT GE
%token LET LETREC IN CASE OF SEMICOLON LAMBDA DOT COMMA ARROW
%token EOF

%type <Syntax.core_sc_defn list> program
%start program

%right prec_let SEMICOLON ARROW
%left EQ NE LT LE GT GE
%left ADD SUB
%left MUL DIV
%left LPAREN
%left INT IDENT PACK

%%

program:
  | sc SEMICOLON? EOF { [$1] }
  | sc SEMICOLON program { $1 :: $3 }

sc:
  | idents ASSIGN expr { (List.hd $1, List.tl $1, $3) }

idents:
  | idents IDENT { $1 @ [$2] }
  | IDENT { [$1] }

idents0:
  | idents { $1 }
  | { [] }

expr:
  | expr aexpr { EAp($1, $2) }
  | aexpr { $1 }
  | LET defns IN expr %prec prec_let { ELet(false, $2, $4) }
  | LETREC defns IN expr %prec prec_let { ELet(true, $2, $4) }
  | CASE expr OF alts %prec prec_let { ECase($2, $4) }
  | LAMBDA idents DOT expr %prec prec_let { ELam($2, $4) }
  | expr MUL expr { EAp(EAp(EVar "*", $1), $3) }
  | expr DIV expr { EAp(EAp(EVar "/", $1), $3) }
  | expr ADD expr { EAp(EAp(EVar "+", $1), $3) }
  | expr SUB expr { EAp(EAp(EVar "-", $1), $3) }
  | expr EQ expr { EAp(EAp(EVar "==", $1), $3) }
  | expr NE expr { EAp(EAp(EVar "!=", $1), $3) }
  | expr LT expr { EAp(EAp(EVar "<", $1), $3) }
  | expr LE expr { EAp(EAp(EVar "<=", $1), $3) }
  | expr GT expr { EAp(EAp(EVar ">", $1), $3) }
  | expr GE expr { EAp(EAp(EVar ">=", $1), $3) }

aexpr:
  | IDENT { EVar($1) }
  | INT { ENum($1) }
  | LPAREN expr RPAREN { $2 }
  | PACK LBRACKET INT COMMA INT RBRACKET { EConstr($3, $5) }

defns:
  | defns SEMICOLON defn { $1 @ [$3] }
  | defn { [$1] }

defn:
  | IDENT ASSIGN expr { ($1, $3) }

alts:
  | alts SEMICOLON alt { $1 @ [$3] }
  | alt { [$1] }

alt:
  | L INT R idents0 ARROW expr { ($2, $4, $6) }

pprint.ml

open Syntax

let (%) f g x = f (g x)

module LazyList = struct
  type 'a node = Nil | Cons of 'a * 'a t
  and 'a t = 'a node Lazy.t
  let empty = lazy Nil
  let singleton x = lazy (Cons (x, empty))
  let force = Lazy.force
  let rec map f l = lazy (
    match force l with
    | Nil -> Nil
    | Cons (h, t) -> Cons (f h, map f t)
  )
  let rec append l1 l2 = lazy (
    match force l1 with
    | Nil -> force l2
    | Cons (h, t) -> Cons (h, append t l2)
  )
  let rec concat ll = lazy (
    match force ll with
    | Nil -> Nil
    | Cons (h, t) -> append h (concat t) |> force
  )
  let is_empty l = force l = Nil
end

let force = Lazy.force

type doc =
  | Empty
  | Line of bool
  | Nest of int * doc
  | Char of char
  | Text of string
  | Cat of doc * doc
  | Union of doc * doc
  | Column of (int -> doc)
  | Nesting of (int -> doc)

type sdoc =
  | SEmpty
  | SChar of char * sdoc Lazy.t
  | SText of string * sdoc Lazy.t
  | SLine of int * sdoc Lazy.t

let rec flatten = function
  | Empty -> Empty
  | Line false -> Empty
  | Line true -> Text " "
  | Nest (i, x) -> flatten x
  | Char c -> Char c
  | Text s -> Text s
  | Cat (x, y) -> Cat (flatten x, flatten y)
  | Union (x, y) -> flatten x
  | Column f -> Column (flatten % f)
  | Nesting f -> Nesting (flatten % f)

let group x = Union (flatten x, x)
let (<|>) x y = Union (x, y)
let (<.>) x y = Cat (x, y)
let space = Text " "
let (<+>) x y = Cat (x, Cat (space, y))
let (</>) x y = Cat (x, Cat (group (Line true), y))
let (<+/>) x y = Cat (x, Cat (Union (space, Line true), y))
let (<$>) x y = Cat (x, Cat (Line true, y))
let align x = Column (fun k -> Nesting (fun i -> Nest (k-i, x)))
let softline b = group (Line b)
let int i = Text (string_of_int i)
let intw w i =
  let s = string_of_int i in
  let l = String.length s in
  if w > l then
    Text (String.make (w-l) ' ' ^ s)
  else
    Text s
let width x f = Column (fun k1 -> x <.> Column (fun k2 -> f (k2-k1)))
let vsep = List.fold_left (<$>)

let fill_break f x =
  width x (fun w ->
    if w > f then Nest (f, Line false)
    else Text (String.make (f-w) ' '))

let fill f x =
  width x (fun w ->
    if w >= f then Empty
    else Text (String.make (f-w) ' '))

let rec sep_by sep = function
  | [] -> Empty
  | x::xs ->
    let rec go acc = function
      | [] -> acc
      | x::xs -> go (acc <.> sep <.> x) xs
    in
    align (go x xs)

let enclose l r x = l <.> x <.> r

let rec enclose_sep l r sep = function
  | [] -> l <.> r
  | x::xs ->
    let rec go acc = function
      | [] -> acc
      | x::xs -> go (acc <.> sep <.> x) xs
    in
    align (go (l <.> x) xs <.> r)

let render rfrac w x =
  let r = rfrac *. float_of_int w |> int_of_float |> min w |> max 0 in
  let better n k x y =
    let rec fits w = function
      | _ when w < 0 -> false
      | SEmpty -> true
      | SChar (c, x) -> fits (w-1) (force x)
      | SText (s, x) -> fits (w-String.length s) (force x)
      | SLine (i, x) -> true
    in
    if fits (min (w-k) (r-k+n)) x then x else y
  in
  let rec best n k = function
    | [] -> SEmpty
    | (i,d)::ds ->
        match d with
        | Empty -> best n k ds
        | Line _ -> SLine (i, lazy (best i i ds))
        | Nest (j,x) -> best n k ((i+j,x)::ds)
        | Char c -> SChar (c, lazy (best n (k+1) ds))
        | Text s -> SText (s, lazy (best n (k+String.length s) ds))
        | Cat (x,y) -> best n k ((i,x)::(i,y)::ds)
        | Union (x,y) -> better n k (best n k ((i,x)::ds)) (best n k ((i,y)::ds))
        | Column f -> best n k ((i,f k)::ds)
        | Nesting f -> best n k ((i,f i)::ds)
  in
  best 0 0 [0,x]

let layout rfrac w x =
  let buf = Buffer.create 0 in
  let rec go = function
    | SEmpty ->
        Buffer.contents buf
    | SChar (c, x) ->
        Buffer.add_char buf c;
        go (force x)
    | SText (s, x) ->
        Buffer.add_string buf s;
        go (force x)
    | SLine (i, x) ->
        Buffer.add_char buf '\n';
        Buffer.add_string buf (String.make i ' ');
        go (force x)
  in
  go (render rfrac w x)

let pp_vars xs =
  sep_by (Char ' ') (List.map (fun x -> Text x) xs)

let paren b x =
  if b then Char '(' <.> x <.> Char ')' else x

let rec pp_def (name, e) =
  Text name <+/> Char '=' <+/> pp_expr 0 e
and pp_defs defs =
  sep_by (Char ';' <.> softline true) (List.map pp_def defs)
and pp_alt (i, vars, e) =
  Char '<' <.> int i <.> Char '>' <.>
  (if vars = [] then Empty else Char ' ' <.> pp_vars vars) <+>
  Text "->" <.>
  group (Nest (2, Line true)) <.> pp_expr 1 e
and pp_alts alts =
  sep_by (Char ';' <.> Line true) (List.map pp_alt alts)
and pp_expr prec = function
  | ENum i ->
      int i
  | EConstr (tag,n) ->
      Text "Pack{" <.> int tag <.> Char ',' <.> int n <.> Char '}'
  | EVar v ->
      Text v
  | EAp (EAp (EVar "*", e1), e2) ->
      paren (prec > 4) (pp_expr 4 e1 <+/> Char '*' <+/> pp_expr 5 e2) |> align |> group
  | EAp (EAp (EVar "/", e1), e2) ->
      paren (prec > 4) (pp_expr 4 e1 <+/> Char '/' <+/> pp_expr 5 e2) |> align |> group
  | EAp (EAp (EVar "+", e1), e2) ->
      paren (prec > 2) (pp_expr 2 e1 <+/> Char '+' <+/> pp_expr 3 e2) |> align |> group
  | EAp (EAp (EVar "-", e1), e2) ->
      paren (prec > 2) (pp_expr 2 e1 <+/> Char '-' <+/> pp_expr 3 e2) |> align |> group
  | EAp (e1, e2) ->
      paren (prec > 6) (pp_expr 6 e1 <+/> pp_expr 7 e2) |> align |> group
  | ELet (isrec, defs, e) ->
      let keyword = Text (if isrec then "letrec" else "let") in
      let ppe = pp_expr 0 e in
      let r = keyword <.>
        group (Nest (2, Line true <.> (pp_defs defs)) <$> Text "in") <$>
        pp_expr 0 e
      in
      paren (prec > 0) (align r) |> group
  | ECase (e, alts) ->
      let r = Text "case" <+/> pp_expr 0 e <+/> Text "of" <.>
        Nest (2, Line true <.> pp_alts alts)
      in
      paren (prec > 0) (align r)
  | ELam (vars, e) ->
      let r = Char '\\' <+/> pp_vars vars <.> Char '.' <+/>
        pp_expr 0 e
      in
      paren (prec > 0) (align r)

let pp_sc (name, params, e) =
  let lhs = pp_vars (name::params) <+/> Char '=' in
  match e with
  | ELet _ ->
      lhs <.> Nest (2, Line true <.> pp_expr 0 e)
  | _ ->
    lhs <+/> align (pp_expr 0 e)

let pp_program scs =
  sep_by (Char ';' <.> Line true <.> Line false) (List.map pp_sc scs) <.> Line false

let pp_list xs =
  sep_by (Line true) @@ List.mapi (fun i x -> intw 3 i <.> Char ')' <+> x) xs

syntax.ml

type name = string

type 'a expr =
  | EVar of name
  | ENum of int
  | EConstr of int * int
  | EAp of 'a expr * 'a expr
  | ELet of bool * ('a * 'a expr) list * 'a expr
  | ECase of 'a expr * 'a alter list
  | ELam of 'a list * 'a expr
and 'a alter = int * 'a list * 'a expr

type 'a sc_defn = name * 'a list * 'a expr

type core_expr = name expr
type core_alt = name alter
type core_sc_defn = name sc_defn
type core_program = core_sc_defn list

% ocamlbuild main.byte
% cat in
take i xs = if (or (i == 0) (empty xs)) nil (cons (head xs) (take (i-1) (tail xs)));
add x y = x + y;
zipwith f xs ys = cons (f (head xs) (head ys)) (zipwith f (tail xs) (tail ys));
fibs = cons 0 (cons 1 (zipwith add fibs (tail fibs)));

main = printList (take 12 fibs);
% ./main.byte in 2>/dev/null
0
1
1
2
3
5
8
13
21
34
55
89