implemented edge and vertex attributes

lib/dot_of_tast.ml

@@ -20,30 +20,46 @@ module L = struct
   let pop () = match !stack with [] -> () | _ :: xs -> stack := xs
 end
 
-let vertex_attributes = ref (fun _ -> [])
-let edge_attributes = ref (fun _ -> [])
-let graph_name = ref None
+open Graph
 
-module G = Graph.Imperative.Digraph.ConcreteBidirectional (struct
-  type t = string
+module G = struct
+  module EDGE = struct
+    type t = string * Graphviz.DotAttributes.edge list
 
-  let compare = String.compare
-  let equal = String.equal
-  let hash = Hashtbl.hash
-end)
+    let compare (a, _) (b, _) = String.compare a b
+    let default = ("", [])
+  end
 
-module Dot = Graph.Graphviz.Dot (struct
-  include G
+  module VERTEX = struct
+    type t = string * Graphviz.DotAttributes.vertex list
 
-  let edge_attributes k = !edge_attributes k
-  let default_edge_attributes _ = []
-  let vertex_name k = k
-  let vertex_attributes k = !vertex_attributes k
+    let compare (a, _) (b, _) = String.compare a b
+    let hash (a, _) = Hashtbl.hash a
+    let equal (a, _) (b, _) = String.equal a b
+  end
+
+  module G = Imperative.Digraph.ConcreteLabeled (VERTEX) (EDGE)
+
+  let vertex_name v =
+    "\"" ^ String.escaped (fst @@ G.V.label v) ^ "\""
+
+  let graph_attributes _ = []
   let default_vertex_attributes _ = []
-  let get_subgraph _ = None
+  let vertex_attributes (_label, attr) = attr
+  let default_edge_attributes _ = []
 
-  let graph_attributes _ =
-    match !graph_name with None -> [] | Some n -> [ `Label n ]
+  let edge_attributes (_src, (label, attr), _dst) =
+    if String.equal label "" then attr else `Label label :: attr
+
+  let get_subgraph (_ : G.V.t) :
+      Graphviz.DotAttributes.subgraph option =
+    None
+
+  include G
+end
+
+module Dot = Graphviz.Dot (struct
+  include G
 end)
 
 let str_of_exp (exp : Typedtree.expression) =
@@ -68,7 +84,7 @@ let str_of_primary_value (exp : Typedtree.expression) =
   | Texp_ifthenelse (_, _, None) -> Fmt.str "if/then"
   | _ -> "????"
 
-let dot_of_impl (s : Typedtree.structure) : unit =
+let graph_of_impl (s : Typedtree.structure) : G.t =
   let g = G.create () in
   let iterator =
     {
@@ -78,7 +94,7 @@ let dot_of_impl (s : Typedtree.structure) : unit =
           match exp.exp_desc with
           | Texp_ident _ | Texp_constant _ | Texp_apply _
           | Texp_function _ | Texp_let _ | Texp_ifthenelse _ -> (
-              let src = F.str "\"%s\"" (str_of_primary_value exp) in
+              let src = str_of_primary_value exp in
               Log.debug (fun m -> m "src: %s" src);
               L.push (F.str "%s" src);
               (match exp.exp_desc with
@@ -93,7 +109,7 @@ let dot_of_impl (s : Typedtree.structure) : unit =
               match L.peek () with
               | Some dst ->
                   Log.debug (fun m -> m "G.add_edge %s %s" src dst);
-                  G.add_edge g src dst
+                  G.add_edge_e g ((src, []), ("", []), (dst, []))
               | None -> ())
           | _ ->
               Tast_iterator.default_iterator.expr iter exp;
@@ -104,9 +120,10 @@ let dot_of_impl (s : Typedtree.structure) : unit =
   in
   iterator.structure iterator s;
   Log.info (fun m ->
-      m "dot_of_impl fprint_graph:@ %a" Dot.fprint_graph g)
+      m "g_of_impl fprint_graph:@ %a" Dot.fprint_graph g);
+  g
 
-let merlin_parse str : unit =
+let merlin_parse str : G.t =
   let config, _command_args =
     Mconfig.parse_arguments ~wd:(Sys.getcwd ())
       ~warning:(fun _ -> ())
@@ -127,5 +144,12 @@ let merlin_parse str : unit =
             (Untypeast.untype_structure t));
       Log.info (fun m ->
           m "Printtyped:@ %a" Ocaml_typing.Printtyped.implementation t);
-      dot_of_impl t
-  | `Interface _ -> ()
+      graph_of_impl t
+  | `Interface _ -> G.create ()
+
+let dot_of_tast ?(fname = "x.dot") str : unit =
+  let oc = open_out fname in
+  F.pf
+    (Format.formatter_of_out_channel oc)
+    "%a" Dot.fprint_graph (merlin_parse str);
+  close_out oc

lib/dot_of_tast.mli

@@ -1 +1 @@
-val merlin_parse : string -> unit
+val dot_of_tast : ?fname:string -> string -> unit