A tree representation for ppx

ppx-tree.ml

open Core_kernel

module Ast = struct
  type ident = string [@@deriving compare, hash, sexp]

  type t =
    | Var of ident
    | Int of int
    | Let of ident * t * t
    | App of ident * t list
  [@@deriving compare, hash, sexp]

  include Base.Comparable.Make(struct
      type nonrec t = t [@@deriving compare, sexp_of]
    end)
end

module Tree = struct
  type tree = {
    uid : int;
    eid : int;
    car : tree;
    cdr : tree;
  }

  type t = tree

  let rec nil = {
    uid = 0;
    eid = 0;
    car = nil;
    cdr = nil;
  }

  module Source = struct
    let last_uid = ref 0
    let last_eid = ref 0

    let trees = Hashtbl.create (module Ast)
    let exps =  Hashtbl.create (module Int)

    let intern ?(car=nil) ?(cdr=nil) exp =
      incr last_eid;
      incr last_uid;
      Hashtbl.add_exn exps ~key:!last_eid ~data:exp;
      {uid = !last_uid; eid = !last_eid; car; cdr}

    let unique t =
      incr last_uid;
      {t with uid = !last_uid}
  end

  let rec refresh t = match t with
    | {uid=0} -> t
    | {car; cdr} -> Source.unique {
        t with
        car = refresh car;
        cdr = refresh cdr
      }

  let rec of_exp exp : tree =
    match Hashtbl.find Source.trees exp with
    | Some t -> refresh t
    | None ->
      let tree = match exp with
        | Var _ | Int _ | App (_,[]) -> Source.intern exp
        | Let (_,x,y) ->
          Source.intern ~car:(of_exp x) ~cdr:(of_exp y) exp
        | App (_,x :: xs) ->
          Source.intern ~car:(of_exp x) ~cdr:(of_exps xs) exp in
      Hashtbl.set Source.trees ~key:exp ~data:tree;
      tree
  and of_exps = function
    | [] -> nil
    | x :: xs -> Source.intern ~car:(of_exp x) ~cdr:(of_exps xs) x

  let to_exp tree = Hashtbl.find_exn Source.exps tree.eid
end