Colors for overlapping binders.

1 files changed, 10 insertions, 3 deletions

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;