diff options
| -rw-r--r-- | discocaml/ast.ml | 41 | ||||
| -rw-r--r-- | discocaml/draw_tree.ml | 11 |
2 files changed, 52 insertions, 0 deletions
diff --git a/discocaml/ast.ml b/discocaml/ast.ml index a4e2272..92f368f 100644 --- a/discocaml/ast.ml +++ b/discocaml/ast.ml @@ -266,3 +266,44 @@ let show_expr (expr : expr ast) : string = Pprintast.expression fmt (parsetree_of_expr expr); Format.pp_print_flush fmt (); Buffer.contents buf + +type binders = expr index option array + +let get_binders (ast : expr ast) : binders = + let out = Array.make (Arraylist.length ast.subexprs) None in + let rec loop (env : (string * expr index) list) (i : expr index) : unit = + match get_subexpr ast i with + | App (f, x) -> + loop env f; + loop env x + | Cons (hd, tl) -> + loop env hd; + loop env tl + | If (cond, then_, else_) -> + loop env cond; + loop env then_; + loop env else_ + | Lam (x, b) -> loop ((x, i) :: env) b + | Let (false, name, bound, body) -> + loop env bound; + loop ((name, i) :: env) body + | Let (true, name, bound, body) -> + loop ((name, i) :: env) bound; + loop ((name, i) :: env) body + | Prim (Add, (l, r)) -> + loop env l; + loop env r + | Prim (Sub, (l, r)) -> + loop env l; + loop env r + | Prim (Mul, (l, r)) -> + loop env l; + loop env r + | Prim (RelOp, (_, l, r)) -> + loop env l; + loop env r + | Var name -> out.(i.index) <- List.assoc_opt name env + | Bool _ | Int _ | Nil -> () + in + loop [] ast.root; + out diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml index 16719df..f0badd7 100644 --- a/discocaml/draw_tree.ml +++ b/discocaml/draw_tree.ml @@ -88,13 +88,24 @@ let draw_tree (ast : expr ast) : string = let fmt = Format.formatter_of_buffer buf in Format.fprintf fmt "digraph {\n"; Format.fprintf fmt " node [shape=\"box\", style=\"rounded\"];\n"; + add_expr_edges ast fmt nodes ast.root; Format.fprintf fmt "\n"; + IntSet.iter (fun index -> let i = { index } in add_node fmt i (get_subexpr ast i)) !nodes; + + get_binders ast + |> Array.iteri (fun i -> function + | Some j -> + Format.fprintf fmt + " expr%d -> expr%d_var [color=\"#cccccc\", constraint=false];\n" i + j.index + | None -> ()); + Format.fprintf fmt "}\n"; Format.pp_print_flush fmt (); Buffer.contents buf |