From 401bc1c5485e26329598d59a140e51f70f58857c Mon Sep 17 00:00:00 2001
From: Nathan Ringo <nathan@remexre.com>
Date: Fri, 19 Jan 2024 17:33:46 -0600
Subject: Fix diagrams, allow multiple roles.

---
 discocaml/ast.ml       | 10 ++++++++++
 discocaml/draw_tree.ml | 22 +++++++++++-----------
 2 files changed, 21 insertions(+), 11 deletions(-)

(limited to 'discocaml')

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index b5ebba3..a706a19 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -142,6 +142,16 @@ let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
           (Pexp_apply
              ( Wrap.expression (Pexp_ident (Wrap.ident "+")),
                [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
+    | Prim (`Sub, [| l; r |]) ->
+        Wrap.expression
+          (Pexp_apply
+             ( Wrap.expression (Pexp_ident (Wrap.ident "-")),
+               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
+    | Prim (`Mul, [| l; r |]) ->
+        Wrap.expression
+          (Pexp_apply
+             ( Wrap.expression (Pexp_ident (Wrap.ident "*")),
+               [ (Nolabel, loop (subexpr l)); (Nolabel, loop (subexpr r)) ] ))
     | Prim (p, xs) ->
         failwith
           ("illegal Prim: " ^ [%derive.show: prim * expr index array] (p, xs))
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index fe3f037..6ea4ff2 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -2,17 +2,17 @@ open Ast
 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 -> 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
+  match expr with
+  | App _ ->
+      Format.fprintf fmt
+        "  expr%d [fontname=\"CMU Typewriter Text Bold\", label=\"apply\"];\n"
+        i.index
+  | Int n -> Format.fprintf fmt "  expr%d [label=\"%d\"];\n" i.index n
+  | Lam (_, _) -> Format.fprintf fmt "  expr%d [label=\"λ\"];\n" i.index
+  | Prim (`Add, _) -> Format.fprintf fmt "  expr%d [label=\"+\"];\n" i.index
+  | Prim (`Sub, _) -> Format.fprintf fmt "  expr%d [label=\"-\"];\n" i.index
+  | Prim (`Mul, _) -> Format.fprintf fmt "  expr%d [label=\"*\"];\n" i.index
+  | Var n -> Format.fprintf fmt "  expr%d [label=%S];\n" i.index n
 
 let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
     (nodes : IntSet.t ref) : expr index -> unit =
-- 
cgit v1.2.3