diff options
Diffstat (limited to 'discocaml')
| -rw-r--r-- | discocaml/arraylist.ml | 1 | ||||
| -rw-r--r-- | discocaml/arraylist.mli | 1 | ||||
| -rw-r--r-- | discocaml/ast.ml | 13 | ||||
| -rw-r--r-- | discocaml/draw_tree.ml | 70 |
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; (* |