initial

lib/dot.ml

@@ -0,0 +1,143 @@
+let escape s =
+  let b = Buffer.create (String.length s) in
+  String.iter
+    (fun c ->
+      match c with
+      | '"' -> Buffer.add_string b "\\\""
+      | '\\' -> Buffer.add_string b "\\\\"
+      | '\n' -> Buffer.add_string b "\\n"
+      | c -> Buffer.add_char b c)
+    s;
+  Buffer.contents b
+
+let sanitize s =
+  String.map
+    (fun c ->
+      if
+        (c >= 'a' && c <= 'z')
+        || (c >= 'A' && c <= 'Z')
+        || (c >= '0' && c <= '9')
+        || c = '_'
+      then c
+      else '_')
+    s
+
+let nt_id name = "nt_" ^ sanitize name
+
+let prob_of probs lhs rhs =
+  match Hashtbl.find_opt probs lhs with
+  | None -> 0.0
+  | Some l -> (
+      match List.find_opt (fun (q, _) -> Grammar.equal_rhs q.Grammar.rhs rhs) l with
+      | Some (_, x) -> x
+      | None -> 0.0)
+
+let prod_key p =
+  (p.Grammar.lhs, String.concat " " (List.map Grammar.symbol_to_string p.Grammar.rhs))
+
+let dot_grammar ?(highlight_nts = []) ?(highlight_prods = [])
+    ?(title = "grammar") g =
+  let probs = Grammar.probabilities g in
+  let hl_nt name = List.mem name highlight_nts in
+  let hl_prods = List.map prod_key highlight_prods in
+  let buf = Buffer.create 2048 in
+  Buffer.add_string buf "digraph grammar {\n";
+  Buffer.add_string buf "  rankdir=LR;\n";
+  Buffer.add_string buf "  labelloc=\"t\";\n";
+  Buffer.add_string buf (Printf.sprintf "  label=\"%s\";\n" (escape title));
+  Buffer.add_string buf
+    "  node [fontname=\"Helvetica\", fontsize=10];\n  edge [fontname=\"Helvetica\", fontsize=9];\n";
+  List.iter
+    (fun nt ->
+      let is_start = String.equal nt g.Grammar.start in
+      let hl = hl_nt nt in
+      Buffer.add_string buf
+        (Printf.sprintf
+           "  %s [label=\"%s\", shape=%s%s%s];\n"
+           (nt_id nt) (escape nt)
+           (if is_start then "doublecircle" else "ellipse")
+           (if hl then ", color=red, penwidth=2.5, fontcolor=red" else "")
+           (if is_start then ", style=bold" else "")))
+    (Grammar.nonterminals g);
+  List.iter
+    (fun t ->
+      Buffer.add_string buf
+        (Printf.sprintf
+           "  tm_%s [label=\"%s\", shape=plaintext, fontcolor=\"#666666\"];\n"
+           (sanitize t) (escape t)))
+    (Grammar.terminals g);
+  List.iteri
+    (fun i p ->
+      let pr = prob_of probs p.Grammar.lhs p.Grammar.rhs in
+      let hl = List.mem (prod_key p) hl_prods in
+      let rhs_text =
+        String.concat " " (List.map Grammar.symbol_to_string p.Grammar.rhs)
+      in
+      Buffer.add_string buf
+        (Printf.sprintf
+           "  pr_%d [label=\"%s\\np=%.3f\", shape=box%s];\n"
+           i (escape rhs_text) pr
+           (if hl then ", color=red, penwidth=2.5, fontcolor=red" else ""));
+      Buffer.add_string buf
+        (Printf.sprintf "  %s -> pr_%d;\n" (nt_id p.Grammar.lhs) i);
+      List.iteri
+        (fun k s ->
+          let target =
+            match s with
+            | Grammar.Nonterm x -> nt_id x
+            | Grammar.Term t -> "tm_" ^ sanitize t
+          in
+          Buffer.add_string buf
+            (Printf.sprintf "  pr_%d -> %s [label=\"%d\"];\n" i target (k + 1)))
+        p.Grammar.rhs)
+    g.Grammar.productions;
+  Buffer.add_string buf "}\n";
+  Buffer.contents buf
+
+let dot_tree ?(title = "parse") tree =
+  let buf = Buffer.create 1024 in
+  Buffer.add_string buf "digraph parse {\n";
+  Buffer.add_string buf "  rankdir=TB;\n";
+  Buffer.add_string buf (Printf.sprintf "  label=\"%s\";\n" (escape title));
+  Buffer.add_string buf
+    "  node [fontname=\"Helvetica\", fontsize=10, shape=ellipse];\n";
+  let counter = ref 0 in
+  let rec go tree =
+    let id = !counter in
+    incr counter;
+    match tree with
+    | Parse.Leaf s ->
+        Buffer.add_string buf
+          (Printf.sprintf "  n%d [label=\"%s\", shape=plaintext, fontcolor=\"#666666\"];\n"
+             id (escape s));
+        id
+    | Parse.Node (nt, children) ->
+        Buffer.add_string buf
+          (Printf.sprintf "  n%d [label=\"%s\"];\n" id (escape nt));
+        List.iteri
+          (fun k child ->
+            let cid = go child in
+            Buffer.add_string buf
+              (Printf.sprintf "  n%d -> n%d [label=\"%d\"];\n" id cid (k + 1)))
+          children;
+        id
+  in
+  ignore (go tree);
+  Buffer.add_string buf "}\n";
+  Buffer.contents buf
+
+let write_file path contents =
+  let oc = open_out path in
+  output_string oc contents;
+  close_out oc
+
+let render ?out_svg ~out_dot dot_string =
+  write_file out_dot dot_string;
+  match out_svg with
+  | None -> None
+  | Some svg ->
+      let cmd =
+        Printf.sprintf "dot -Tsvg %s -o %s 2>/dev/null"
+          (Filename.quote out_dot) (Filename.quote svg)
+      in
+      if Sys.command cmd = 0 then Some svg else None