Colors for overlapping binders.

4 files changed, 72 insertions, 13 deletions

diff --git a/discocaml/arraylist.ml b/discocaml/arraylist.ml
index 110cd17..8979349 100644
--- a/discocaml/arraylist.ml
+++ b/discocaml/arraylist.ml
@@ -56,6 +56,7 @@ let push arraylist x =
   arraylist.len <- arraylist.len + 1
 
 let to_array { array; len; _ } = Array.sub array 0 len
+let to_seq arraylist = Array.to_seq (to_array arraylist)
 
 module Array_for_pp = struct
   type 'a t = 'a array [@@deriving show { with_path = false }]
diff --git a/discocaml/arraylist.mli b/discocaml/arraylist.mli
index 6b40ce0..cf42000 100644
--- a/discocaml/arraylist.mli
+++ b/discocaml/arraylist.mli
@@ -7,4 +7,5 @@ val get : 'a t -> int -> 'a
 val set : 'a t -> int -> 'a -> unit
 val push : 'a t -> 'a -> unit
 val to_array : 'a t -> 'a array
+val to_seq : 'a t -> 'a Seq.t
 val pp : (Format.formatter -> 'a -> unit) -> Format.formatter -> 'a t -> unit
diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index 92f368f..86630a6 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -267,7 +267,8 @@ let show_expr (expr : expr ast) : string =
   Format.pp_print_flush fmt ();
   Buffer.contents buf
 
-type binders = expr index option array
+type binder_info = [ `Binder of string | `Bound of expr index ]
+type binders = binder_info option array
 
 let get_binders (ast : expr ast) : binders =
   let out = Array.make (Arraylist.length ast.subexprs) None in
@@ -283,11 +284,15 @@ let get_binders (ast : expr ast) : binders =
         loop env cond;
         loop env then_;
         loop env else_
-    | Lam (x, b) -> loop ((x, i) :: env) b
+    | Lam (x, b) ->
+        out.(i.index) <- Some (`Binder x);
+        loop ((x, i) :: env) b
     | Let (false, name, bound, body) ->
+        out.(i.index) <- Some (`Binder name);
         loop env bound;
         loop ((name, i) :: env) body
     | Let (true, name, bound, body) ->
+        out.(i.index) <- Some (`Binder name);
         loop ((name, i) :: env) bound;
         loop ((name, i) :: env) body
     | Prim (Add, (l, r)) ->
@@ -302,7 +307,9 @@ let get_binders (ast : expr ast) : binders =
     | Prim (RelOp, (_, l, r)) ->
         loop env l;
         loop env r
-    | Var name -> out.(i.index) <- List.assoc_opt name env
+    | Var name ->
+        out.(i.index) <-
+          List.assoc_opt name env |> Option.map (fun j -> `Bound j)
     | Bool _ | Int _ | Nil -> ()
   in
   loop [] ast.root;
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index 14c08af..60a8aa5 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -8,7 +8,53 @@ let fmt_with_possible_subscript (fmt : Format.formatter) (s : string) : unit =
         sub
   | None -> Format.fprintf fmt "%S" s
 
-let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
+let colors =
+  [| "#56b4e9"; "#009e73"; "#e69f00"; "#d55e00"; "#cc79a7"; "#ffaa14" |]
+
+let get_color (i : int) = colors.(i mod Array.length colors)
+
+let assign_colors (ast : expr ast) : string array =
+  let binders = get_binders ast in
+  let out = Array.make (Arraylist.length ast.subexprs) "#000000" in
+
+  (* First, accumulate lists of binders for each "basename". *)
+  let name_counts = Hashtbl.create 16 in
+  binders
+  |> Array.iteri (fun i -> function
+       | Some (`Binder name) ->
+           let base =
+             match Util.break_to_subscript name with
+             | Some (base, _) -> base
+             | None -> name
+           in
+           let binders =
+             match Hashtbl.find_opt name_counts base with
+             | Some binders -> binders
+             | None ->
+                 let binders = Arraylist.make 0 0 in
+                 Hashtbl.add name_counts base binders;
+                 binders
+           in
+           Arraylist.push binders i
+       | _ -> ());
+
+  (* Next, go through the binders to find those with basename collisions. If
+     there's overlap, assign colors to the binders. *)
+  name_counts
+  |> Hashtbl.iter (fun _ binders ->
+         if Arraylist.length binders > 1 then
+           Arraylist.to_seq binders
+           |> Seq.iteri (fun i j -> out.(j) <- get_color i));
+
+  (* Finally, find the bound variables and copy their colors. *)
+  binders
+  |> Array.iteri (fun i -> function
+       | Some (`Bound j) -> out.(i) <- out.(j.index) | _ -> ());
+
+  out
+
+let add_node (fmt : Format.formatter) (i : expr index) (expr : expr)
+    (color : string) : unit =
   match expr with
   | App _ ->
       Format.fprintf fmt
@@ -44,11 +90,11 @@ let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
       Format.fprintf fmt "  expr%d [label=\"%s\"];\n" i.index
         (string_of_relop op)
   | Var x ->
-      Format.fprintf fmt "  expr%d [label=%a];\n" i.index
+      Format.fprintf fmt "  expr%d [color=\"%s\", label=%a];\n" i.index color
         fmt_with_possible_subscript x
 
 let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
-    (nodes : IntSet.t ref) : expr index -> unit =
+    (nodes : IntSet.t ref) (colors : string array) : unit =
   let rec loop (i : expr index) : unit =
     nodes := IntSet.add i.index !nodes;
     let edge_to (j : expr index) : unit =
@@ -68,13 +114,15 @@ let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
         edge_to else_
     | Lam (x, b) ->
         Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
-        Format.fprintf fmt "  expr%d_var [label=%a, shape=\"ellipse\"];\n"
-          i.index fmt_with_possible_subscript x;
+        Format.fprintf fmt
+          "  expr%d_var [color=\"%s\", label=%a, shape=\"ellipse\"];\n" i.index
+          colors.(i.index) fmt_with_possible_subscript x;
         edge_to b
     | Let (_, name, bound, body) ->
         Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
-        Format.fprintf fmt "  expr%d_var [label=%a, shape=\"ellipse\"];\n"
-          i.index fmt_with_possible_subscript name;
+        Format.fprintf fmt
+          "  expr%d_var [color=\"%s\", label=%a, shape=\"ellipse\"];\n" i.index
+          colors.(i.index) fmt_with_possible_subscript name;
         edge_to bound;
         edge_to body
     | Prim (Add, (l, r)) ->
@@ -92,7 +140,7 @@ let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
     | Var _ -> ()
     | Bool _ | Int _ | Nil -> ()
   in
-  loop
+  loop ast.root
 
 let draw_tree (ast : expr ast) : string =
   let buf = Buffer.create 16 and nodes = ref IntSet.empty in
@@ -102,13 +150,15 @@ let draw_tree (ast : expr ast) : string =
     "  node [fontsize=\"20pt\", penwidth=\"2.0\", shape=\"box\", \
      style=\"rounded\"];\n";
 
-  add_expr_edges ast fmt nodes ast.root;
+  let colors = assign_colors ast in
+
+  add_expr_edges ast fmt nodes colors;
   Format.fprintf fmt "\n";
 
   IntSet.iter
     (fun index ->
       let i = { index } in
-      add_node fmt i (get_subexpr ast i))
+      add_node fmt i (get_subexpr ast i) colors.(i.index))
     !nodes;
 
   (*