closer

lib/dot_of_tast.ml

@@ -59,7 +59,11 @@ let _ =
   Logs.set_reporter (Logs_fmt.reporter ());
   Logs.set_level (Some Debug)
 
-module F = Fmt
+module F = struct
+  include Fmt
+
+  let bar ppf () = string ppf "|"
+end
 
 module Log = (val Logs.src_log
                     (Logs.Src.create "dot_of_tast"
@@ -704,7 +708,16 @@ let str_of_exp (exp : Typedtree.expression) =
       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 pp_value_binding_list ppf (vbl : Typedtree.value_binding list) =
   Fmt.pf ppf "%a"
@@ -716,11 +729,12 @@ let unique_node_count = ref 0
 
 let vertex_of_primary_value (exp : Typedtree.expression) : Gr.Vertex.t
     =
+  (* sorry that the use of the word primary is so abusive *)
   match exp.exp_desc with
   | Texp_ident (path, _longident, _value_desc) ->
       {
         name = F.str "%a" P.fmt_path path;
-        attr = [ `Label (F.str "{{<0>}|<out>%s}" (str_of_exp exp)) ];
+        attr = [ `Label (F.str "{{}|<out>%s}" (str_of_exp exp)) ];
         subgraph = None;
       }
   | Texp_constant _ ->
@@ -732,105 +746,154 @@ let vertex_of_primary_value (exp : Typedtree.expression) : Gr.Vertex.t
   | Texp_variant _ -> { Gr.Vertex.default with name = str_of_exp exp }
   | Texp_apply (exp, args) ->
       let cnt = ref 0 in
-      {
+      let label =
-        Gr.Vertex.default with
+        F.str "{{%a}|<out>%s}"
-        name = str_of_exp exp;
+          (F.list ~sep:F.bar (fun ppf -> function
-        attr =
-          [
-            `Label
-              (F.str "{{%a}|<o>%s}"
-                 (F.list
-                    ~sep:(fun ppf () -> F.string ppf "|")
-                    (fun ppf arg_list ->
-                      match arg_list with
              | 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));
+          args (str_of_exp exp)
-          ];
+      in
-      }
-  | Texp_function { arg_label = _; param; cases = _; partial = _ } ->
       {
         Gr.Vertex.default with
-        name = Fmt.str "fun %s" (str_of_ident param);
+        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 (_, _, Some _) ->
+  | Texp_ifthenelse (_, _, _) -> Gr.Vertex.default
-      {
-        Gr.Vertex.default with
-        (* TODO make unique (with loc?) *)
-        name = "if/then/else";
-        attr = [ `Label "{{<0>if|<1>then|<2>else}|<o>}" ];
-      }
-  | Texp_ifthenelse (_, _, None) ->
-      {
-        Gr.Vertex.default with
-        name = Fmt.str "if/then";
-        attr = [ `Label "{{<1>if|<2>then}|<o>}" ];
-      }
   | _ ->
       {
         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 sdst : (Gr.V.t * int) Stack.t = Stack.create () in
+  let dst_stk : (Gr.V.t * int) Stack.t = Stack.create () in
-  let iterator =
+  let expr (iter : Tast_iterator.iterator)
-    {
+      (exp : Typedtree.expression) =
-      Tast_iterator.default_iterator with
+    let extra = function
-      expr =
+      | Typedtree.Texp_constraint cty -> iter.typ iter cty
-        (fun iter exp ->
+      | Texp_coerce (cty1, cty2) ->
-          let src = vertex_of_primary_value exp in
+          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
-          Log.debug (fun m -> m "src: %s" (Gr.Vertex.label src));
     incr unique_node_count;
-          (match exp.exp_desc with
+    let src =
+      let open Typedtree in
+      match exp.exp_desc with
       | Texp_apply (_exp, list) ->
+          let src = vertex_of_primary_value 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) sdst;
+                  Stack.push (src, i) dst_stk;
                   iter.expr iter ee;
-                      Stack.pop sdst |> ignore)
+                  Stack.pop dst_stk |> ignore)
                 expopt;
               match arg_label with
               | Asttypes.Nolabel -> i + 1
               | _ -> i)
             0 list
-              |> ignore
+          |> ignore;
+          src
       | Texp_ifthenelse (e0, e1, e2) ->
-              Stack.push (src, 0) sdst;
+          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 sdst |> ignore;
+          Stack.pop dst_stk |> ignore;
-              Stack.push (src, 1) sdst;
+          Stack.push (src, 1) dst_stk;
           iter.expr iter e1;
-              Stack.pop sdst |> ignore;
+          Stack.pop dst_stk |> ignore;
           Option.iter
             (fun e' ->
-                  Stack.push (src, 2) sdst;
+              Stack.push (src, 2) dst_stk;
               iter.expr iter e';
-                  Stack.pop sdst |> ignore)
+              Stack.pop dst_stk |> ignore)
-                e2
+            e2;
+          src
+      | Texp_function { arg_label = _; param; cases; partial = _ } ->
+          let src =
+            labeled_vertex
+              (unique_str_of_ident param)
+              (F.str "{{<0>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 (str_of_pat c_lhs) "{<in>|LOL}" );
+              Stack.pop_opt dst_stk |> ignore)
+            cases;
+          src
       | _ ->
-              Stack.push (src, 0) sdst;
+          let src = vertex_of_primary_value exp in
+          Stack.push (src, 0) dst_stk;
           Tast_iterator.default_iterator.expr iter exp;
-              Stack.pop sdst |> ignore);
+          Stack.pop_opt dst_stk |> ignore;
-          match Stack.top_opt sdst with
+          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)
+            m "G.add_edge %s:%s %s:%s" (Gr.Vertex.label src) srcport
-                    srcport (Gr.Vertex.label dst) (F.str "%d" dstport));
+              (Gr.Vertex.label dst) (F.str "%d" dstport));
         Gr.add_edge_e g
           ( src,
             Gr.Edge.
@@ -843,12 +906,10 @@ let graph_of_impl (s : Typedtree.structure) : Gr.t =
                   ];
               },
             dst )
-          | None -> ());
+    | None -> ()
-    }
   in
+  let iterator = { Tast_iterator.default_iterator with expr } in
   iterator.structure iterator s;
-  Log.info (fun m ->
-      m "g_of_impl Printtyped:@ %a" Printtyped.implementation s);
   g
 
 let merlin_parse str : Gr.t =
@@ -869,7 +930,9 @@ let merlin_parse str : Gr.t =
   | `Implementation t ->
       Log.info (fun m ->
           m "Untyped:@ %a" Pprintast.structure
-            (Untypeast.untype_structure t));
+            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