From 752d36d855946909ddba30176fb6ee14eadfb761 Mon Sep 17 00:00:00 2001
From: tali <iitalics@gmail.com>
Date: Wed, 13 Dec 2023 18:16:04 -0500
Subject: [PATCH] deny referencing variables outside of captured scope

---
 bin/main.ml        |  7 +++++-
 lib/compile/bcc.ml |  5 +++--
 lib/compile/ir.ml  | 56 +++++++++++++++++++++++++++++-----------------
 3 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/bin/main.ml b/bin/main.ml
index 1029a36..e680df0 100644
--- a/bin/main.ml
+++ b/bin/main.ml
@@ -5,7 +5,12 @@ let () =
   Logs.set_level (Some Logs.Debug);
 
   try
-    let ast = parse "val two = 2  fun twice(x) x*two  val r = twice(8)" in
+    let ast = parse "
+                val two = 2
+                val zero = 0
+                fun twice(x) x*two+zero
+                println(twice(4))
+ " in
     Logs.debug (fun m -> m "[AST] %a" Ast.pp_modl ast);
     let prog = compile ast in
     Logs.debug (fun m -> Code.dump (m "[BC] %s") prog.main);
diff --git a/lib/compile/bcc.ml b/lib/compile/bcc.ml
index e9a28c7..de4c4a3 100644
--- a/lib/compile/bcc.ml
+++ b/lib/compile/bcc.ml
@@ -22,12 +22,13 @@ let rec compile_lambda (lam : Ir.lambda) =
 
   let reg_of_id = Hashtbl.create 128 in
   let set_reg id r =
-    if Hashtbl.mem reg_of_id id then failwith "BUG: id reassigned";
+    if Hashtbl.mem reg_of_id id then
+      Fmt.failwith "BUG: '%a' reassigned" Ir.pp_id id;
     Hashtbl.add reg_of_id id r
   in
   let get_reg id =
     try Hashtbl.find reg_of_id id with
-      Not_found -> failwith "BUG: id unassigned"
+      Not_found -> Fmt.failwith "BUG: '%a' unassigned" Ir.pp_id id
   in
 
   let rec emit_exp sp = function
diff --git a/lib/compile/ir.ml b/lib/compile/ir.ml
index 66aa6fd..b5e98d9 100644
--- a/lib/compile/ir.ml
+++ b/lib/compile/ir.ml
@@ -7,7 +7,17 @@ let compile_error f =
   Fmt.kstr (fun msg -> raise (Error msg)) f
 
 type imm = Value.t
-type id = Id of int [@@unboxed]
+
+type id = Id of int * string
+(* type id = Id of int [@@unboxed] *)
+
+let make_id_dispenser () =
+  let i = ref (-1) in
+  fun x -> (incr i; Id (!i, x))
+  (* fun _ -> (incr i; Id !i) *)
+
+let pp_id ppf (Id (n, x)) = Fmt.pf ppf "%s_%d" x n
+(* let pp_id ppf (Id n) = Fmt.pf ppf "_%d" n *)
 
 type uop =
   | Not
@@ -52,14 +62,11 @@ and lambda = {
 
 (* lower *)
 
-let make_id_dispenser () =
-  let i = ref (-1) in fun () -> (incr i; Id !i)
-
 module Env = struct
   type t =
     | Obj of {
         self : id;
-        elems : string list;
+        elems : (string, string) Hashtbl.t;
       }
     | Fun of {
         args : (string * id) list;
@@ -73,9 +80,7 @@ module Env = struct
       List.assoc name args, None
 
     | Obj { self; elems } ->
-      if not (List.mem name elems) then
-        raise Not_found;
-      self, Some name
+      self, Some (Hashtbl.find elems name)
 
     | Cons (e1, e2) ->
       try find name e1 with
@@ -142,7 +147,7 @@ let lower ~lib (modl : Ast.modl) =
     match path with
     | Ast.Ele (obj, fld) ->
       let rhs = lower_exp env obj in
-      let lhs = new_id () in
+      let lhs = new_id "get" in
       Let (lhs, rhs, k (`Get (lhs, fld)))
 
     | Ast.Var name ->
@@ -150,26 +155,32 @@ let lower ~lib (modl : Ast.modl) =
       | id, None -> k (`Var id)
       | obj, Some fld -> k (`Get (obj, fld))
       | exception Not_found ->
-        compile_error "unbound variable %S" name
+        compile_error "%S not in scope" name
 
   and lower_block ?(is_scope = false) env items =
     (* collect names of bindings to form the new environment; also check if a scope ends
        with an expression, if not then it is an error *)
-    let elems, ends_with_exp =
+    let elems = Hashtbl.create 32 in
+    let ends_with_exp =
       List.fold_left
-        (fun (elems, _) -> function
-           | Ast.Item_exp _ -> elems, true
+        (fun _ -> function
+           | Ast.Item_exp _ -> true
            | Ast.Item_val (name, _)
            | Ast.Item_obj (name, _)
-           | Ast.Item_fun (name, _, _) -> name :: elems, false)
-        ([], false)
+           | Ast.Item_fun (name, _, _) ->
+             if Hashtbl.mem elems name then
+               compile_error "multiple definitions of %S" name;
+             Hashtbl.add elems name name;
+             false)
+        false
         items
     in
     if is_scope && not ends_with_exp then
       compile_error "scope does not end in expression";
 
-    let self = new_id () in
-    let env = Env.Cons (Obj { self; elems }, env) in
+    let self = new_id "obj" in
+    let env' = Env.Obj { self; elems } in
+    let env = Env.Cons (env', env) in
 
     let funs_r, vals_r, inits_r =
       List.fold_left
@@ -190,7 +201,7 @@ let lower ~lib (modl : Ast.modl) =
              fns, name :: vls, init :: ins
 
            | Ast.Item_fun (name, args, body) ->
-             let fn = (name, lower_lambda self env args body) in
+             let fn = (name, lower_lambda self env' args body) in
              fn :: fns, vls, ins)
         ([], [], [])
         items
@@ -216,7 +227,7 @@ let lower ~lib (modl : Ast.modl) =
     )
 
   and lower_lambda self env args body =
-    let args = List.map (fun a -> a, new_id ()) args in
+    let args = List.map (fun a -> a, new_id a) args in
     let env = Env.Cons (Fun { args }, env) in
     (* TODO: closure conversion *)
     let body = lower_exp env body in
@@ -224,8 +235,11 @@ let lower ~lib (modl : Ast.modl) =
     { self; args; body }
 
   in
-  let self = new_id () in
-  let env = Env.Obj { self; elems = List.map fst lib } in
+
+  let self = new_id "lib" in
+  let elems = Hashtbl.create 128 in
+  List.iter (fun (name, _) -> Hashtbl.add elems name name) lib;
+  let env = Env.Obj { self; elems } in
   let args = [] in
   let body = lower_block env modl.items in
   { self; args; body }