open Ast
module IntSet = Set.Make (Int)

let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
  let label =
    match expr with
    | App _ -> "\"$\""
    | Int n -> Format.sprintf "\"%d\"" n
    | Lam (_, _) -> "\"λ\""
    | Prim (`Add, _) -> "\"+\""
    | Prim (`Sub, _) -> "\"-\""
    | Prim (`Mul, _) -> "\"*\""
    | Var n -> Format.sprintf "%S" n
  in
  Format.fprintf fmt "  expr%d [label=%s];\n" i.index label

let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
    (nodes : IntSet.t ref) : expr index -> unit =
  let rec loop (i : expr index) : unit =
    nodes := IntSet.add i.index !nodes;
    let edge_to (j : expr index) : unit =
      loop j;
      Format.fprintf fmt "  expr%d -> expr%d;\n" i.index j.index
    in
    match get_subexpr ast i with
    | App (f, x) ->
        edge_to f;
        edge_to x
    | Int _ -> ()
    | Lam (x, b) ->
        Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
        Format.fprintf fmt "  expr%d_var [label=%S];\n" i.index x;
        edge_to b
    | Prim (_, xs) -> Array.iter edge_to xs
    | Var _ -> ()
  in
  loop

let draw_tree (ast : expr ast) : string =
  let buf = Buffer.create 16 and nodes = ref IntSet.empty in
  let fmt = Format.formatter_of_buffer buf in
  Format.fprintf fmt "digraph {\n";
  add_expr_edges ast fmt nodes ast.root;
  Format.fprintf fmt "\n";
  IntSet.iter
    (fun index ->
      let i = { index } in
      add_node fmt i (get_subexpr ast i))
    !nodes;
  Format.fprintf fmt "}\n";
  Format.pp_print_flush fmt ();
  Buffer.contents buf