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 _ -> [])
+open Graph
-let edge_attributes = ref (fun _ -> [])
-let graph_name = ref None
 
-module G = Graph.Imperative.Digraph.ConcreteBidirectional (struct
+module G = struct
-  type t = string
+  module EDGE = struct
+    type t = string * Graphviz.DotAttributes.edge list
 
-  let compare = String.compare
+    let compare (a, _) (b, _) = String.compare a b
-  let equal = String.equal
+    let default = ("", [])
-  let hash = Hashtbl.hash
+  end
-end)
 
-module Dot = Graph.Graphviz.Dot (struct
+  module VERTEX = struct
-  include G
+    type t = string * Graphviz.DotAttributes.vertex list
 
-  let edge_attributes k = !edge_attributes k
+    let compare (a, _) (b, _) = String.compare a b
-  let default_edge_attributes _ = []
+    let hash (a, _) = Hashtbl.hash a
-  let vertex_name k = k
+    let equal (a, _) (b, _) = String.equal a b
-  let vertex_attributes k = !vertex_attributes k
+  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 _ =
+  let edge_attributes (_src, (label, attr), _dst) =
-    match !graph_name with None -> [] | Some n -> [ `Label n ]
+    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) =
@@ -78,7 +94,7 @@ let graph_of_impl (s : Typedtree.structure) : G.t =
           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 graph_of_impl (s : Typedtree.structure) : G.t =
               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;