Basic scope construction

1 files changed, 30 insertions, 15 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 78882c1..7f3e0af 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -112,13 +112,15 @@ let rec compile_term env : term -> compile_term_result =
       | Null     -> Application (Identifier "Rego.Any.null", []))
      |> result_of
   | Var id ->
-     (match Map.find_exn env id with
-      | { lean_id; complete = true } -> Identifier lean_id |> result_of
-      | { lean_id; complete = false } ->
+     (match Map.find env id with
+      | Some { lean_id; complete = true } -> Identifier lean_id |> result_of
+      | Some { lean_id; complete = false } ->
         let anonymous_id = gensym "anonymous" in
         { form       = Identifier anonymous_id;
           toplevel   = [];
-          new_locals = [(anonymous_id, lean_id)]; })
+          new_locals = [(anonymous_id, lean_id)]; }
+      | None ->
+        failwith (Printf.sprintf "unbound variable %s" id))
   | Ref (target, terms) ->
     List.fold
       terms
@@ -169,8 +171,14 @@ let rec compile_term env : term -> compile_term_result =
       toplevel = toplevel head @ (List.map ~f:toplevel subterms |> List.concat);
       new_locals = new_locals head @ (List.map ~f:new_locals subterms |> List.concat) }
 
-let rec compile_expr env : expr -> compile_term_result = function
-  | Term t -> compile_term env t
+let rec compile_expr env : expr -> (environ * compile_term_result) = function
+  | Term t ->
+    (match t with
+     (* Special case of call; this is the only kind of term that affects scope. *)
+     | Call ((Var "assign", []), [ Var name; rhs ]) ->
+       let env' = Map.set env ~key:name ~data:{ lean_id = name; complete = true } in
+       (env', compile_term env' (Call ((Var "eq", []), [ Var name; rhs ])))
+     | t -> (env, compile_term env t))
   | Every (Some key, value, domain, query) ->
     (match compile_term env domain with
      | { form = domain'; toplevel; new_locals } ->
@@ -183,7 +191,7 @@ let rec compile_expr env : expr -> compile_term_result = function
                     ((Identifier "Membership.mem"),
                      [Tuple [Identifier key; Identifier value]; domain']));
                  query_compiled])))
-       in {form; toplevel = toplevel @ toplevel'; new_locals})
+       in (env, {form; toplevel = toplevel @ toplevel'; new_locals}))
   | Every (None, value, domain, query) ->
     (match compile_term env domain with
      | { form = domain'; toplevel; new_locals } ->
@@ -196,17 +204,24 @@ let rec compile_expr env : expr -> compile_term_result = function
                     ((Identifier "Membership.mem"),
                      [Identifier value; domain']));
                  query_compiled])))
-       in {form; toplevel = toplevel @ toplevel'; new_locals})
+       in (env, {form; toplevel = toplevel @ toplevel'; new_locals}))
 
-and compile_literal env : literal -> compile_term_result = function
+and compile_literal env : literal -> (environ * compile_term_result) = function
   | Exists (exprs, _) ->
     let (toplevel, form) = compile_query env (List.map ~f:literal_of_expr exprs) in
-    { form; toplevel; new_locals = [] }
+    (env, { form; toplevel; new_locals = [] })
   | Expr (expr, _) -> compile_expr env expr
   | Negate (_expr, _) -> failwith "unimplemented"
 
 and compile_query env literals : (lean_expression list * lean_expression) =
-  let subterms = List.map ~f:(compile_literal env) literals in
+  let (free, bound) = extract_variables env literals in
+  let env' = List.map (free @ bound) ~f:(fun name -> (name, { lean_id = name; complete = true }))
+             |> Environ.of_alist_exn
+             |> Map.merge env ~f:(fun ~key:_key -> function
+                 | `Right v -> Some v
+                 | `Left v -> Some v
+                 | `Both (_, v) -> Some v) in
+  let (_env', subterms) = map_with_environment_updates env' literals compile_literal in
   let (compiled_terms, toplevel, new_locals) =
     List.fold subterms
       ~init:([], [], [])
@@ -224,7 +239,7 @@ and compile_query env literals : (lean_expression list * lean_expression) =
      | [x] -> x
      | xs -> Application (Identifier "And", xs)
    else
-     Exists ((List.map ~f:fst new_locals, "Rego.Any"),
+     Exists ((free @ bound @ (List.map ~f:fst new_locals), "Rego.Any"),
              Application (Identifier "And",
                           List.map ~f:(fun (x, y) ->
                               Application (Identifier y, [Identifier x]))
@@ -251,11 +266,11 @@ let compile_definition env (rules : rule list) =
        List.mapi rules ~f:(fun i ->
            function
            | { head = (_, Assignment (_, term)); body; _ } ->
-             (match compile_term Environ.empty term with
+             (match compile_term toplevel term with
               (* TODO: We need to pass this up.. not ignore it... *)
               | { form; _ } ->
                 (Printf.sprintf "rule%d" i,
-                 compile_query Environ.empty body |> snd,
+                 compile_query toplevel body |> snd,
                  [form]))
            | _ -> failwith "unimplemented"))
   in (env, form)
@@ -283,7 +298,7 @@ let compile_policy env (rules : rule list) =
 
 let compile_package = function
   | { path; imports; policy } ->
-    let env = Environ.empty in
+    let env = toplevel in
     let (_env', forms) = compile_policy env policy in
     [Comment (show_reference path |> String.escaped);
      Comment (List.map ~f:show_import imports |> String.concat |> String.escaped)]