Working diagramming.

2 files changed, 33 insertions, 17 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index d1652d8..b5ebba3 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -13,7 +13,7 @@ type 'a index = { index : int } [@@deriving show { with_path = false }]
 
 let get (arr : 'a Arraylist.t) (i : 'a index) : 'a = Arraylist.get arr i.index
 
-type prim = [ `Add ] [@@deriving show { with_path = false }]
+type prim = [ `Add | `Mul | `Sub ] [@@deriving show { with_path = false }]
 
 type expr =
   | App of expr index * expr index
@@ -37,12 +37,29 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
     Arraylist.push ast.subexprs expr;
     { index }
   in
+  let binop (prim : prim) : expr index =
+    add
+      (Lam
+         ( "$0",
+           add
+             (Lam ("$1", add (Prim (prim, [| add (Var "$0"); add (Var "$1") |]))))
+         ))
+  in
+
   let rec loop (expr : Parsetree.expression) : expr index =
     match expr.pexp_desc with
     | Pexp_apply
         ( { pexp_desc = Pexp_ident { txt = Lident "+"; _ }; _ },
           [ (Nolabel, l); (Nolabel, r) ] ) ->
         add (Prim (`Add, [| loop l; loop r |]))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "-"; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (`Sub, [| loop l; loop r |]))
+    | Pexp_apply
+        ( { pexp_desc = Pexp_ident { txt = Lident "*"; _ }; _ },
+          [ (Nolabel, l); (Nolabel, r) ] ) ->
+        add (Prim (`Mul, [| loop l; loop r |]))
     | Pexp_apply (f, xs) ->
         List.fold_left
           (fun f -> function
@@ -54,15 +71,9 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
         match arg_pat.ppat_desc with
         | Ppat_var { txt = name; _ } -> add (Lam (name, loop body))
         | _ -> raise_unsupported_expr expr)
-    | Pexp_ident { txt = Lident "+"; _ } ->
-        add
-          (Lam
-             ( "$0",
-               add
-                 (Lam
-                    ( "$1",
-                      add (Prim (`Add, [| add (Var "$0"); add (Var "$1") |])) ))
-             ))
+    | Pexp_ident { txt = Lident "+"; _ } -> binop `Add
+    | Pexp_ident { txt = Lident "-"; _ } -> binop `Sub
+    | Pexp_ident { txt = Lident "*"; _ } -> binop `Mul
     | Pexp_ident { txt = Lident name; _ } -> add (Var name)
     | _ -> raise_unsupported_expr expr
   in
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index 30c466f..fe3f037 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -4,13 +4,15 @@ module IntSet = Set.Make (Int)
 let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
   let label =
     match expr with
-    | App _ -> "$"
-    | Int n -> string_of_int n
-    | Lam (x, _) -> "λ " ^ x ^ " →"
-    | Prim (`Add, _) -> "+"
-    | Var n -> n
+    | App _ -> "\"$\""
+    | Int n -> Format.sprintf "\"%d\"" n
+    | Lam (_, _) -> "\"λ\""
+    | Prim (`Add, _) -> "\"+\""
+    | Prim (`Sub, _) -> "\"-\""
+    | Prim (`Mul, _) -> "\"*\""
+    | Var n -> Format.sprintf "%S" n
   in
-  Format.fprintf fmt "  expr%d [label=%S];\n" i.index label
+  Format.fprintf fmt "  expr%d [label=%s];\n" i.index label
 
 let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
     (nodes : IntSet.t ref) : expr index -> unit =
@@ -25,7 +27,10 @@ let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
         edge_to f;
         edge_to x
     | Int _ -> ()
-    | Lam (_, b) -> edge_to b
+    | Lam (x, b) ->
+        Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
+        Format.fprintf fmt "  expr%d_var [label=%S];\n" i.index x;
+        edge_to b
     | Prim (_, xs) -> Array.iter edge_to xs
     | Var _ -> ()
   in