(*

   'let' names nets (scoped vars)
   'fun' names ports (argument vars)



digraph {
  compound=true;
  node [shape=record];
  edge [decorate=true];
  rankdir="LR";

  subgraph G {
    label="let rec fact n = if n = 0 then 1. else float n *. fact (n - 1)";

    cluster=true;
    labeljust="l";
    "1.";
    "if/then/else" [label="{{<1>if|<2>then|<3>else}|<o>}"];
    "1";
    "0";

    minus [label="{{<1>|<2>}|<o>-}"];

  float_mult [label="{{<1>|<2>}|<o>*.}"];
  "float" [label="{<1>|<o>float}"];
  "(=)" [label="{{<1>|<2>}|<o>=}"];
  "fun n" [label="{<o>fun|<1>n}"];
  "factorial" [label="{<1>|<o>factorial}"];
  
  
  "1." -> "if/then/else":2 [label="float"];
  "fun n":1 -> minus:1 [lhead="cluster_minus" label="int"];
  "fun n":1 -> "(=)":2 [label="int"];
  "fun n":1 -> "float":1 [label="int"];
  "if/then/else" -> "->":o [label="float"];
  "1" -> minus:2 [lhead="cluster_minus" label="int"];
  "0" -> "(=)":1 [label="int"];
  minus:o -> "factorial":1 [label="int"];
  float_mult:o -> "if/then/else":3 [label="float"];
  "float":o -> float_mult:1 [label="float"];
  "(=)":o -> "if/then/else":1 [label="bool"];
  "factorial":o -> float_mult:2 [label="float"];
  
  }
}


   
 *)

open Merlin_kernel
open Merlin_utils
open Ocaml_parsing
open Ocaml_typing

let _ =
  Logs.set_reporter (Logs_fmt.reporter ());
  Logs.set_level (Some Debug)

module F = struct
  include Fmt

  let bar ppf () = string ppf "|"
end

module Log = (val Logs.src_log
                    (Logs.Src.create "dot_of_tast"
                       ~doc:"dot_of_tast.ml logger") : Logs.LOG)

open Graph

module DotAttributes = struct
  include Graphviz.DotAttributes

  (* the dot spec probably added records after ocamlgraph/src/graphviz.ml was written *)
  type edge =
    [ Graphviz.DotAttributes.edge
    | `HeadportRecord of string
    | `TailportRecord of string ]
end

module Gr = struct
  open DotAttributes

  module Vertex = struct
    type t = {
      name : string;
      attr : DotAttributes.vertex list;
      subgraph : DotAttributes.subgraph option;
    }

    let compare a b = String.compare a.name b.name
    let hash { name; _ } = Hashtbl.hash name
    let equal a b = compare a b = 0
    let label a = a.name
    let default = { name = ""; attr = []; subgraph = None }
  end

  module Edge = struct
    type t = { name : string; attr : DotAttributes.edge list }

    let compare a b = String.compare a.name b.name
    let default = { name = ""; attr = [] }
  end

  module G = Imperative.Digraph.ConcreteLabeled (Vertex) (Edge)
  module B = Builder.I (G)
  module V = G.V
  module E = G.E

  type t = G.t * graph list

  let iter_vertex a (b, _) = G.iter_vertex a b
  let iter_edges_e a (b, _) = G.iter_edges_e a b
  let add_edge_e (a, _) b = G.add_edge_e a b
  let vertex_name v = F.str "\"%s\"" (String.escaped (Vertex.label v))
  let graph_attributes (_g, attr) = attr
  let default_vertex_attributes _ = [ `Shape `Record ]
  let vertex_attributes (a : V.t) = a.attr
  let default_edge_attributes _ = [ `Decorate true ]

  let edge_attributes ((_src, { name; attr; _ }, _dst) : E.t) =
    if String.equal name "" then attr else `Label name :: attr

  let get_subgraph (v : G.V.t) :
      Graphviz.DotAttributes.subgraph option =
    v.subgraph

  let create ?(size = 32) attr = (B.G.create ~size (), attr)
end

module Dot = struct
  module X = struct
    include Gr
  end

  open Format

  module EN = struct
    type color = int
    type color_with_transparency = int32

    let color_to_color_with_transparency color =
      Int32.add (Int32.shift_left (Int32.of_int color) 8) 0xFFl

    let fprint_color ppf color = fprintf ppf "\"#%06X\"" color

    let fprint_color_with_transparency ppf color =
      fprintf ppf "\"#%08lX\"" color

    let fprint_string ppf s = fprintf ppf "\"%s\"" s
    (* let s' = String.escaped s in
       if s' = s && s <> ""
       then fprintf ppf "%s" s
       else fprintf ppf "\"%s\"" s'*)

    let fprint_string_user ppf s =
      (*  let s = String.escaped s in*)
      fprintf ppf "\"%s\"" s

    let fprint_htmlstring_user ppf s = fprintf ppf "<%s>" s

    let fprint_square_not_empty printer ppf = function
      | [] -> ()
      | l -> fprintf ppf " [%a]" printer l

    type arrow_style =
      [ `None
      | `Normal
      | `Onormal
      | `Inv
      | `Dot
      | `Odot
      | `Invdot
      | `Invodot ]

    let fprint_arrow_style ppf = function
      | `None -> fprintf ppf "none"
      | `Normal -> fprintf ppf "normal"
      | `Onormal -> fprintf ppf "onormal"
      | `Inv -> fprintf ppf "inv"
      | `Dot -> fprintf ppf "dot"
      | `Odot -> fprintf ppf "odot"
      | `Invdot -> fprintf ppf "invdot"
      | `Invodot -> fprintf ppf "invodot"

    let fprint_dir ppf = function
      | `TopToBottom -> fprintf ppf "TB"
      | `BottomToTop -> fprintf ppf "BT"
      | `LeftToRight -> fprintf ppf "LR"
      | `RightToLeft -> fprintf ppf "RL"

    type symbseq = COMMA | SEMI

    let fprint_symbseq ppf = function
      | COMMA -> pp_print_string ppf ","
      | SEMI -> pp_print_string ppf ";"

    module Attributes = struct
      module CommonAttributes = struct
        include Graphviz.CommonAttributes
        (** Pretty-print. *)

        let fprint_orientation ppf = function
          | `Portrait -> fprintf ppf "portrait"
          | `Landscape -> fprintf ppf "landscape"

        let fprint_graph ppf = function
          | `Center b -> fprintf ppf "center=%i" (if b then 1 else 0)
          | `Fontcolor a -> fprintf ppf "fontcolor=%a" fprint_color a
          | `Fontname s -> fprintf ppf "fontname=%a" fprint_string s
          | `Fontsize i -> fprintf ppf "fontsize=%i" i
          | `Label s -> fprintf ppf "label=%a" fprint_string_user s
          | `HtmlLabel s ->
              fprintf ppf "label=%a" fprint_htmlstring_user s
          | `Orientation a ->
              fprintf ppf "orientation=%a" fprint_orientation a
          | `Page (x, y) -> fprintf ppf "page=\"%f,%f\"" x y
          | `Pagedir a -> fprintf ppf "pagedir=%a" fprint_dir a
          | `Size (x, y) -> fprintf ppf "size=\"%f,%f\"" x y
          | `OrderingOut -> fprintf ppf "ordering=out"

        let fprint_shape ppf = function
          | `Ellipse -> fprintf ppf "ellipse"
          | `Box -> fprintf ppf "box"
          | `Circle -> fprintf ppf "circle"
          | `Doublecircle -> fprintf ppf "doublecircle"
          | `Diamond -> fprintf ppf "diamond"
          | `Plaintext -> fprintf ppf "plaintext"
          | `Record -> fprintf ppf "record"
          | `Egg -> fprintf ppf "egg"
          | `House -> fprintf ppf "house"
          | `Invhouse -> fprintf ppf "invhouse"
          | `Trapezium -> fprintf ppf "trapezium"
          | `Invtrapezium -> fprintf ppf "invtrapezium"
          | `Triangle -> fprintf ppf "triangle"
          | `Invtriangle -> fprintf ppf "invtriangle"
          | `Oval -> fprintf ppf "oval"
          | `Assembly -> fprintf ppf "assembly"
          | `Box3d -> fprintf ppf "box3d"
          | `Cds -> fprintf ppf "cds"
          | `Component -> fprintf ppf "component"
          | `Doubleoctagon -> fprintf ppf "doubleoctagon"
          | `Fivepoverhang -> fprintf ppf "fivepoverhang"
          | `Folder -> fprintf ppf "folder"
          | `Insulator -> fprintf ppf "insulator"
          | `Larrow -> fprintf ppf "larrow"
          | `Lpromoter -> fprintf ppf "lpromoter"
          | `Mcircle -> fprintf ppf "mcircle"
          | `Mdiamond -> fprintf ppf "mdiamond"
          | `Msquare -> fprintf ppf " msquare"
          | `Note -> fprintf ppf "note"
          | `Noverhang -> fprintf ppf "noverhang"
          | `Parallelogram -> fprintf ppf "parallelogram"
          | `Primersite -> fprintf ppf "primersite"
          | `Promoter -> fprintf ppf "promoter"
          | `Proteasesite -> fprintf ppf "proteasesite"
          | `Proteinstab -> fprintf ppf "proteinstab"
          | `Rarrow -> fprintf ppf "rarrow"
          | `Restrictionsite -> fprintf ppf "restrictionsite"
          | `Ribosite -> fprintf ppf "ribosite"
          | `Rnastab -> fprintf ppf "rnastab"
          | `Rpromoter -> fprintf ppf "rpromoter"
          | `Signature -> fprintf ppf "signature"
          | `Star -> fprintf ppf "star"
          | `Tab -> fprintf ppf "tab"
          | `Terminator -> fprintf ppf "terminator"
          | `Threepoverhang -> fprintf ppf "threepoverhang"
          | `Tripleoctagon -> fprintf ppf "tripleoctagon"
          | `Underline -> fprintf ppf "underline"
          | `Utr -> fprintf ppf "utr"
          | `Polygon (i, f) ->
              fprintf ppf "polygon, sides=%i, skew=%f" i f

        let rec fprint_string_list ppf = function
          | [] -> ()
          | [ hd ] -> fprintf ppf "%s" hd
          | hd :: tl -> fprintf ppf "%s,%a" hd fprint_string_list tl

        let node_style_str = function
          | `Rounded -> "rounded"
          | `Filled -> "filled"
          | `Solid -> "solid"
          | `Dashed -> "dashed"
          | `Dotted -> "dotted"
          | `Bold -> "bold"
          | `Invis -> "invis"

        let fprint_style_list sep ppf a =
          fprintf ppf "style=\"%a\"%a@ " fprint_string_list
            (List.map node_style_str a)
            fprint_symbseq sep

        let fprint_vertex ppf = function
          | `Color a -> fprintf ppf "color=%a" fprint_color a
          | `ColorWithTransparency a ->
              fprintf ppf "color=%a" fprint_color_with_transparency a
          | `Fontcolor a -> fprintf ppf "fontcolor=%a" fprint_color a
          | `Fontname s -> fprintf ppf "fontname=%a" fprint_string s
          | `Fontsize i -> fprintf ppf "fontsize=%i" i
          | `Height f -> fprintf ppf "height=%f" f
          | `Label s -> fprintf ppf "label=%a" fprint_string_user s
          | `HtmlLabel s ->
              fprintf ppf "label=%a" fprint_htmlstring_user s
          | `Orientation f -> fprintf ppf "orientation=%f" f
          | `Penwidth f -> fprintf ppf "penwidth=%f" f
          | `Peripheries i -> fprintf ppf "peripheries=%i" i
          | `Regular b -> fprintf ppf "regular=%b" b
          | `Shape a -> fprintf ppf "shape=%a" fprint_shape a
          | `Style _ -> assert false
          | `Width f -> fprintf ppf "width=%f" f

        let fprint_arrow_direction ppf = function
          | `Forward -> fprintf ppf "forward"
          | `Back -> fprintf ppf "back"
          | `Both -> fprintf ppf "both"
          | `None -> fprintf ppf "none"

        let fprint_edge ppf = function
          | `Color a -> fprintf ppf "color=%a" fprint_color a
          | `ColorWithTransparency a ->
              fprintf ppf "color=%a" fprint_color_with_transparency a
          | `Decorate b -> fprintf ppf "decorate=%b" b
          | `Dir a -> fprintf ppf "dir=%a" fprint_arrow_direction a
          | `Fontcolor a -> fprintf ppf "fontcolor=%a" fprint_color a
          | `Fontname s -> fprintf ppf "fontname=%a" fprint_string s
          | `Fontsize i -> fprintf ppf "fontsize=%i" i
          | `Label s -> fprintf ppf "label=%a" fprint_string_user s
          | `HtmlLabel s ->
              fprintf ppf "label=%a" fprint_htmlstring_user s
          | `Labelfontcolor a ->
              fprintf ppf "labelfontcolor=%a" fprint_color a
          | `Labelfontname s -> fprintf ppf "labelfontname=\"%s\"" s
          (* (String.escaped s) *)
          | `Labelfontsize i -> fprintf ppf "labelfontsize=%i" i
          | `Penwidth f -> fprintf ppf "penwidth=%f" f
          | `Style _ -> assert false

        let rec filter_style al sl l =
          match l with
          | [] -> (al, sl)
          | `Style s :: l -> filter_style al (s :: sl) l
          | a :: l -> filter_style (a :: al) sl l

        (** [fprint_graph_attribute printer ppf list] pretty prints a list of
      attributes on the formatter [ppf], using the printer [printer] for
      each attribute.  The list appears between brackets and attributes
      are speparated by ",".  If the list is empty, nothing is printed. *)
        let fprint_attributes fprint_style_list fprint_attribute sep
            ppf list =
          if list <> [] then (
            let list, styles = filter_style [] [] list in
            let rec fprint_attributes_rec ppf = function
              | [] -> ()
              | hd :: tl ->
                  fprintf ppf "%a%a@ " fprint_attribute hd
                    fprint_symbseq sep;
                  fprint_attributes_rec ppf tl
            in
            fprintf ppf "@[<hov>%a" fprint_attributes_rec list;
            if styles <> [] then fprint_style_list sep ppf styles;
            fprintf ppf "@]")
      end

      include DotAttributes
      (** {4 Pretty-print of attributes} *)

      let rec fprint_string_list ppf = function
        | [] -> ()
        | [ hd ] -> fprintf ppf "%s" hd
        | hd :: tl -> fprintf ppf "%s,%a" hd fprint_string_list tl

      let fprint_ratio ppf = function
        | `Float f -> fprintf ppf "%f" f
        | `Fill -> fprintf ppf "fill"
        | `Compress -> fprintf ppf "compress"
        | `Auto -> fprintf ppf "auto"

      let fprint_graph ppf = function
        | #CommonAttributes.graph as att ->
            CommonAttributes.fprint_graph ppf att
        | `Bgcolor a -> fprintf ppf "bgcolor=%a" fprint_color a
        | `BgcolorWithTransparency a ->
            fprintf ppf "bgcolor=%a" fprint_color_with_transparency a
        | `Comment s -> fprintf ppf "comment=%a" fprint_string s
        | `Concentrate b -> fprintf ppf "concentrate=%b" b
        | `Fontpath s -> fprintf ppf "fontpath=%a" fprint_string s
        | `Layers s -> fprintf ppf "layers=%a" fprint_string_list s
        | `Margin f -> fprintf ppf "margin=%f" f
        | `Mclimit f -> fprintf ppf "mclimit=%f" f
        | `Nodesep f -> fprintf ppf "nodesep=%f" f
        | `Nslimit i -> fprintf ppf "nslimit=%i" i
        | `Nslimit1 i -> fprintf ppf "nslimit1=%i" i
        | `Ranksep f -> fprintf ppf "ranksep=%f" f
        | `Quantum f -> fprintf ppf "quantum=%f" f
        | `Rankdir a -> fprintf ppf "rankdir=%a" fprint_dir a
        | `Ratio a -> fprintf ppf "ratio=%a" fprint_ratio a
        | `Samplepoints i -> fprintf ppf "samplepoints=%i" i
        | `Url s -> fprintf ppf "URL=\"%s\"" s (*(String.escaped s)*)

      let fprint_vertex ppf = function
        | #CommonAttributes.vertex as att ->
            CommonAttributes.fprint_vertex ppf att
        | `Comment s -> fprintf ppf "comment=%a" fprint_string s
        | `Distortion f -> fprintf ppf "distortion=%f" f
        | `Fillcolor a -> fprintf ppf "fillcolor=%a" fprint_color a
        | `FillcolorWithTransparency a ->
            fprintf ppf "fillcolor=%a" fprint_color_with_transparency
              a
        | `Fixedsize b -> fprintf ppf "fixedsize=%b" b
        | `Layer s -> fprintf ppf "layer=%a" fprint_string s
        | `Url s -> fprintf ppf "URL=\"%s\"" s (*(String.escaped s)*)
        | `Z f -> fprintf ppf "z=%f" f

      let fprint_port ppf = function
        | `N -> fprintf ppf "n"
        | `NE -> fprintf ppf "ne"
        | `E -> fprintf ppf "e"
        | `SE -> fprintf ppf "se"
        | `S -> fprintf ppf "s"
        | `SW -> fprintf ppf "sw"
        | `W -> fprintf ppf "w"
        | `NW -> fprintf ppf "nw"

      let fprint_edge ppf = function
        | #CommonAttributes.edge as att ->
            CommonAttributes.fprint_edge ppf att
        | `Arrowhead a ->
            fprintf ppf "arrowhead=%a" fprint_arrow_style a
        | `Arrowsize f -> fprintf ppf "arrowsize=%f" f
        | `Arrowtail a ->
            fprintf ppf "arrowtail=%a" fprint_arrow_style a
        | `Comment s -> fprintf ppf "comment=%a" fprint_string s
        | `Constraint b -> fprintf ppf "constraint=%b" b
        | `Headlabel s -> fprintf ppf "headlabel=%a" fprint_string s
        | `Headport a -> fprintf ppf "headport=%a" fprint_port a
        | `HeadportRecord a ->
            fprintf ppf "headport=%a" fprint_string a
        | `Headurl s -> fprintf ppf "headURL=%a" fprint_string s
        | `Labelangle f -> fprintf ppf "labelangle=%f" f
        | `Labeldistance f -> fprintf ppf "labeldistance=%f" f
        | `Labelfloat b -> fprintf ppf "labelfloat=%b" b
        | `Layer s -> fprintf ppf "layer=%a" fprint_string s
        | `Minlen i -> fprintf ppf "minlen=%i" i
        | `Samehead s -> fprintf ppf "samehead=%a" fprint_string s
        | `Sametail s -> fprintf ppf "sametail=%a" fprint_string s
        | `Taillabel s -> fprintf ppf "taillabel=%a" fprint_string s
        | `Tailport a -> fprintf ppf "tailport=%a" fprint_port a
        | `TailportRecord a ->
            fprintf ppf "tailport=%a" fprint_string a
        | `Tailurl s -> fprintf ppf "tailURL=%a" fprint_string s
        | `Weight i -> fprintf ppf "weight=%i" i

      let fprint_vertex_list =
        CommonAttributes.fprint_attributes
          CommonAttributes.fprint_style_list fprint_vertex

      let fprint_edge_list =
        CommonAttributes.fprint_attributes
          CommonAttributes.fprint_style_list fprint_edge
    end

    let name = "dot"
    let opening = "digraph"
    let edge_arrow = "->"
  end

  open Format
  (** {3 Common attributes} *)

  type color = int
  type color_with_transparency = int32

  let color_to_color_with_transparency color =
    Int32.add (Int32.shift_left (Int32.of_int color) 8) 0xFFl

  let fprint_color ppf color = fprintf ppf "\"#%06X\"" color

  let fprint_color_with_transparency ppf color =
    fprintf ppf "\"#%08lX\"" color

  let fprint_string ppf s = fprintf ppf "\"%s\"" s
  (* let s' = String.escaped s in
     if s' = s && s <> ""
     then fprintf ppf "%s" s
     else fprintf ppf "\"%s\"" s'*)

  let fprint_string_user ppf s =
    (*  let s = String.escaped s in*)
    fprintf ppf "\"%s\"" s

  let fprint_htmlstring_user ppf s = fprintf ppf "<%s>" s

  let fprint_square_not_empty printer ppf = function
    | [] -> ()
    | l -> fprintf ppf " [%a]" printer l

  type arrow_style =
    [ `None
    | `Normal
    | `Onormal
    | `Inv
    | `Dot
    | `Odot
    | `Invdot
    | `Invodot ]

  let fprint_arrow_style ppf = function
    | `None -> fprintf ppf "none"
    | `Normal -> fprintf ppf "normal"
    | `Onormal -> fprintf ppf "onormal"
    | `Inv -> fprintf ppf "inv"
    | `Dot -> fprintf ppf "dot"
    | `Odot -> fprintf ppf "odot"
    | `Invdot -> fprintf ppf "invdot"
    | `Invodot -> fprintf ppf "invodot"

  let fprint_dir ppf = function
    | `TopToBottom -> fprintf ppf "TB"
    | `BottomToTop -> fprintf ppf "BT"
    | `LeftToRight -> fprintf ppf "LR"
    | `RightToLeft -> fprintf ppf "RL"

  type symbseq = COMMA | SEMI

  let fprint_symbseq ppf = function
    | COMMA -> pp_print_string ppf ","
    | SEMI -> pp_print_string ppf ";"

  (* had to copy out of ocamlgraph/src/graphviz.ml because people don't make shit extensible *)
  let command = ref EN.name
  let set_command cmd = command := cmd

  exception Error of string

  let handle_error f arg =
    try f arg
    with Error msg ->
      Printf.eprintf "%s: %s failure\n   %s\n" Sys.argv.(0) EN.name
        msg;
      flush stderr;
      exit 2

  (** [fprint_graph_attributes ppf list] pretty prints a list of
      graph attributes on the formatter [ppf].  Attributes are separated
      by a ";". *)
  let fprint_graph_attributes ppf list =
    List.iter
      (function
        | att -> fprintf ppf "%a;@ " EN.Attributes.fprint_graph att)
      list

  (** [fprint_graph ppf graph] pretty prints the graph [graph] in
      the CGL language on the formatter [ppf]. *)
  let fprint_graph ppf graph =
    let module SG = Map.Make (String) in
    let subgraphs = ref SG.empty in

    (* Printing nodes. *)
    let print_nodes ppf =
      let default_node_attributes =
        X.default_vertex_attributes graph
      in
      if default_node_attributes <> [] then
        fprintf ppf "node%a;@ "
          (fprint_square_not_empty
             (EN.Attributes.fprint_vertex_list COMMA))
          default_node_attributes;

      X.iter_vertex
        (function
          | node ->
              (match X.get_subgraph node with
              | None -> ()
              | Some sg ->
                  let sg, nodes =
                    if SG.mem sg.EN.Attributes.sg_name !subgraphs then
                      SG.find sg.EN.Attributes.sg_name !subgraphs
                    else (sg, [])
                  in
                  subgraphs :=
                    SG.add sg.EN.Attributes.sg_name
                      (sg, node :: nodes)
                      !subgraphs);
              fprintf ppf "%s%a;@ " (X.vertex_name node)
                (fprint_square_not_empty
                   (EN.Attributes.fprint_vertex_list COMMA))
                (X.vertex_attributes node))
        graph
    in

    (* Printing subgraphs *)
    let rec print_nested_subgraphs ppf = function
      | [] -> () (* no more work to do, so terminate *)
      | name :: worklist ->
          let sg, nodes = SG.find name !subgraphs in
          let children =
            SG.filter
              (fun _ (sg, _) ->
                sg.EN.Attributes.sg_parent = Some name)
              !subgraphs
          in
          fprintf ppf "@[<v 2>subgraph cluster_%s { %a%t@ %t };@]@\n"
            name
            (EN.Attributes.fprint_vertex_list SEMI)
            sg.EN.Attributes.sg_attributes
            (fun ppf ->
              List.iter
                (fun n -> fprintf ppf "%s;" (X.vertex_name n))
                nodes)
            (fun ppf ->
              print_nested_subgraphs ppf
                (List.map fst (SG.bindings children)));

          print_nested_subgraphs ppf worklist
    in

    let print_subgraphs ppf =
      let root_worklist =
        SG.filter
          (fun _ (sg, _) -> sg.EN.Attributes.sg_parent = None)
          !subgraphs
      in
      print_nested_subgraphs ppf
        (List.map fst (SG.bindings root_worklist))
    in

    (* Printing edges *)
    let print_edges ppf =
      let default_edge_attributes = X.default_edge_attributes graph in
      if default_edge_attributes <> [] then
        fprintf ppf "edge%a;@ "
          (fprint_square_not_empty
             (EN.Attributes.fprint_edge_list COMMA))
          default_edge_attributes;

      X.iter_edges_e
        (function
          | edge ->
              fprintf ppf "%s %s %s%a;@ "
                (X.vertex_name (X.E.src edge))
                EN.edge_arrow
                (X.vertex_name (X.E.dst edge))
                (fprint_square_not_empty
                   (EN.Attributes.fprint_edge_list COMMA))
                (X.edge_attributes edge))
        graph
    in

    fprintf ppf "@[<v>%s G {@ @[<v 2>  %a" EN.opening
      fprint_graph_attributes
      (X.graph_attributes graph);
    fprintf ppf "%t@ " print_nodes;
    fprintf ppf "%t@ " print_subgraphs;
    fprintf ppf "%t@ " print_edges;
    fprintf ppf "@]}@;@]"

  (** [output_graph oc graph] pretty prints the graph [graph] in the dot
      language on the channel [oc]. *)
  let output_graph oc graph =
    let ppf = formatter_of_out_channel oc in
    fprint_graph ppf graph;
    pp_print_flush ppf ()
end

module Pmapper = struct
  type t = Ast_mapper.mapper

  let default = Ast_mapper.default_mapper

  let remove_attributes : t =
    { default with attributes = (fun _m _alist -> []) }
end

(* Stolen from ocaml/typing/printtyped.ml cause it's not exposed!! *)
module P = struct
  open Format
  open Lexing
  open Location
  module Path = Ocaml_typing.Path

  let rec fmt_longident_aux f x =
    match x with
    | Longident.Lident s -> fprintf f "%s" s
    | Longident.Ldot (y, s) -> fprintf f "%a.%s" fmt_longident_aux y s
    | Longident.Lapply (y, z) ->
        fprintf f "%a(%a)" fmt_longident_aux y fmt_longident_aux z

  let fmt_longident f x = fprintf f "\"%a\"" fmt_longident_aux x.txt
  let fmt_ident = Ident.print

  let fmt_modname f = function
    | None -> fprintf f "_"
    | Some id -> Ident.print f id

  let rec fmt_path f x =
    match x with
    | Path.Pident s -> fprintf f "%a" fmt_ident s
    | Path.Pdot (y, s) -> fprintf f "%a.%s" fmt_path y s
    | Path.Papply (y, z) -> fprintf f "%a(%a)" fmt_path y fmt_path z
end

let str_of_exp (exp : Typedtree.expression) =
  Fmt.str "%a" Pprintast.expression
    Pmapper.(
      remove_attributes.expr remove_attributes
        (Untypeast.untype_expression exp))

open Typedtree

let str_of_pat (pat : value general_pattern) =
  Fmt.str "%a" Pprintast.pattern
    Pmapper.(
      remove_attributes.pat remove_attributes
        (Untypeast.untype_pattern pat))

let str_of_ident = Ident.name
let unique_str_of_ident = Ident.unique_name

let unique_str_of_ident_pat x =
  match x.pat_desc with
  | Tpat_var (s, _) -> F.str "%a" P.fmt_ident s
  | _ -> str_of_pat x

let pp_value_binding_list ppf (vbl : Typedtree.value_binding list) =
  Fmt.pf ppf "%a"
    (Fmt.parens
       (Fmt.list (fun ppf (ident, _loc, _ty) -> Ident.print ppf ident)))
    (Typedtree.let_bound_idents_full vbl)

let unique_node_count = ref 0

let vertex_of_exp (exp : Typedtree.expression) : Gr.Vertex.t =
  (* primary isn't well defined here, it's whatever i choosen for src nodes atm *)
  match exp.exp_desc with
  | Texp_ident (path, _longident, _value_desc) ->
      {
        name = F.str "%a" P.fmt_path path;
        attr = [ `Label (F.str "{{<in>}|<out>%s}" (str_of_exp exp)) ];
        subgraph = None;
      }
  | Texp_construct _ ->
      {
        Gr.Vertex.default with
        name = F.str "_Texp_construct/%d" !unique_node_count;
        attr = [ `Label (F.str "{<out>%s}" (str_of_exp exp)) ];
      }
  | Texp_constant _ ->
      {
        Gr.Vertex.default with
        name = F.str "_Texp_constant/%d" !unique_node_count;
        attr = [ `Label (F.str "{<0>|<out>%s}" (str_of_exp exp)) ];
      }
  | Texp_variant _ -> { Gr.Vertex.default with name = str_of_exp exp }
  | Texp_apply (exp, args) ->
      let cnt = ref 0 in
      let label =
        F.str "{{%a}|<out>%s}"
          (F.list ~sep:F.bar (fun ppf -> function
             | Asttypes.Nolabel, _exp ->
                 F.pf ppf "<%d>" !cnt;
                 cnt := !cnt + 1
             | Asttypes.(Labelled s | Optional s), _exp ->
                 F.pf ppf "<%s>" s))
          args (str_of_exp exp)
      in
      {
        Gr.Vertex.default with
        name = str_of_exp exp;
        attr = [ `Label label ];
      }
  | Texp_let (_, vbl, _) ->
      {
        Gr.Vertex.default with
        name = Fmt.str "let %a" pp_value_binding_list vbl;
      }
  | Texp_ifthenelse (_, _, _) -> Gr.Vertex.default
  | _ ->
      {
        Gr.Vertex.default with
        name = Fmt.str "???? (%s)" (str_of_exp exp);
      }

let labeled_vertex name label =
  { Gr.Vertex.default with name; attr = [ `Label label ] }

let graph_of_impl (s : Typedtree.structure) : Gr.t =
  let g = Gr.create [ `Rankdir `LeftToRight ] in
  let dst_stk : (Gr.V.t * int) Stack.t = Stack.create () in
  let expr (iter : Tast_iterator.iterator)
      (exp : Typedtree.expression) =
    let extra = function
      | Typedtree.Texp_constraint cty -> iter.typ iter cty
      | Texp_coerce (cty1, cty2) ->
          Option.iter (iter.typ iter) cty1;
          iter.typ iter cty2
      | Texp_newtype _ -> ()
      | Texp_newtype' _ -> ()
      | Texp_poly cto -> Option.iter (iter.typ iter) cto
    in
    List.iter (fun (e, _, _) -> extra e) exp.exp_extra;
    iter.env iter exp.exp_env;

    let srcport = "out" in
    incr unique_node_count;
    let src =
      let open Typedtree in
      match exp.exp_desc with
      | Texp_apply (_exp, list) ->
          let src = vertex_of_exp exp in
          (* Tast_iterator.default_iterator.expr iter exp *)
          (* skip _exp *)
          List.fold_left
            (fun i (arg_label, expopt) ->
              Option.iter
                (fun ee ->
                  Stack.push (src, i) dst_stk;
                  iter.expr iter ee;
                  Stack.pop dst_stk |> ignore)
                expopt;
              match arg_label with
              | Asttypes.Nolabel -> i + 1
              | _ -> i)
            0 list
          |> ignore;
          src
      | Texp_ifthenelse (e0, e1, e2) ->
          let src =
            match e2 with
            | Some _ ->
                {
                  Gr.Vertex.default with
                  (* TODO make unique (with loc?) *)
                  name = "if/then/else";
                  attr = [ `Label "{{<0>if|<1>then|<2>else}|<out>}" ];
                }
            | None ->
                {
                  Gr.Vertex.default with
                  name = "if/then";
                  attr = [ `Label "{{<1>if|<2>then}|<out>}" ];
                }
          in
          Stack.push (src, 0) dst_stk;
          iter.expr iter e0;
          Stack.pop dst_stk |> ignore;
          Stack.push (src, 1) dst_stk;
          iter.expr iter e1;
          Stack.pop dst_stk |> ignore;
          Option.iter
            (fun e' ->
              Stack.push (src, 2) dst_stk;
              iter.expr iter e';
              Stack.pop dst_stk |> ignore)
            e2;
          src
      | Texp_function { arg_label = _; param; cases; partial = _ } ->
          let src =
            labeled_vertex
              (unique_str_of_ident param)
              (F.str "{{<in>function %s}|{%s}}"
                 (unique_str_of_ident param)
                 (fst
                    (List.fold_left
                       (fun (s, i) { c_lhs; c_guard; _ } ->
                         ( F.str "%s|<%d>%s%s" s i (str_of_pat c_lhs)
                             (Option.fold ~none:""
                                ~some:(fun _ -> " with TODO c_guard")
                                c_guard),
                           i + 1 ))
                       ("", 0) cases))
                 (* -> *))
          in
          List.iteri
            (fun i { c_lhs; c_rhs; _ } ->
              Stack.push (src, i) dst_stk;
              iter.expr iter c_rhs;
              Gr.add_edge_e g
                ( src,
                  Gr.Edge.
                    {
                      name = "";
                      attr =
                        [
                          `HeadportRecord "in";
                          `TailportRecord (F.str "%d" i);
                        ];
                    },
                  labeled_vertex
                    (unique_str_of_ident_pat c_lhs)
                    "{<in>|LOL}" );
              Stack.pop_opt dst_stk |> ignore)
            cases;
          src
      | _ ->
          let src = vertex_of_exp exp in
          Stack.push (src, 0) dst_stk;
          Tast_iterator.default_iterator.expr iter exp;
          Stack.pop_opt dst_stk |> ignore;
          src
    in
    match Stack.top_opt dst_stk with
    | Some (dst, dstport) ->
        Log.debug (fun m ->
            m "G.add_edge %s:%s %s:%s" (Gr.Vertex.label src) srcport
              (Gr.Vertex.label dst) (F.str "%d" dstport));
        Gr.add_edge_e g
          ( src,
            Gr.Edge.
              {
                name = "";
                attr =
                  [
                    `HeadportRecord (F.str "%d" dstport);
                    `TailportRecord srcport;
                  ];
              },
            dst )
    | None -> ()
  in
  let iterator = { Tast_iterator.default_iterator with expr } in
  iterator.structure iterator s;
  g

let merlin_parse str : Gr.t =
  let config, _command_args =
    Mconfig.parse_arguments ~wd:(Sys.getcwd ())
      ~warning:(fun _ -> ())
      (List.map snd []) [] Mconfig.initial []
  in
  File_id.with_cache @@ fun () ->
  let source = Msource.make str in
  let pipeline = Mpipeline.make config source in
  Mpipeline.with_pipeline pipeline @@ fun () ->
  Log.info (fun m -> m "merlin_parse...");
  let typer = Mpipeline.typer_result pipeline in
  let typedtree = Mtyper.get_typedtree typer in
  let _initial_env = Mtyper.initial_env typer in
  let mbrowse = Mbrowse.of_typedtree typedtree in
  Log.info (fun m -> m "Mbrowse.print %s" (Mbrowse.print () mbrowse));
  match typedtree with
  | `Implementation t ->
      Log.info (fun m ->
          m "Untyped:@ %a" Pprintast.structure
            Pmapper.(
              remove_attributes.structure remove_attributes
                (Untypeast.untype_structure t)));
      Log.info (fun m ->
          m "Printtyped:@ %a" Ocaml_typing.Printtyped.implementation t);
      graph_of_impl t
  | `Interface _ -> Gr.create []

let dot_of_tast ?(fname = "x.dot") str : unit =
  let oc = open_out fname in
  Dot.output_graph oc (merlin_parse str);
  close_out oc