Error on variables used before assignment

1 files changed, 16 insertions, 3 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 540b094..83fa856 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -71,8 +71,12 @@ let extract_variables (env : environ) (q : query) : analysis_result =
       else { free = [name]; bound = []; used = [] }
     (* Otherwise, we'll recurse on the term. *)
     | Ref (target, args) ->
-      let terms = [term_of_ref_target target] @ args in
-      List.map terms ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
+      let target = term_of_ref_target target in
+      if List.is_empty args then
+        handle_term env target
+      else
+        let terms = [] @ args in
+        List.map terms ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
     | Array xs | Set xs ->
       List.map xs ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
     | Object kvs -> List.map ~f:snd kvs
@@ -124,6 +128,8 @@ let rec compile_term env : term -> lean_expression =
   | Var id ->
     (match Map.find env id with
      | Some { lean_id; complete = true; assigned = true } -> Identifier lean_id
+     | Some { assigned = false; _ } ->
+       failwith (Printf.sprintf "identifier %s used before definition" id)
      | Some _ -> failwith "unimplemented"
      | None -> failwith (Printf.sprintf "unbound variable %s" id))
   | Ref (target, terms) ->
@@ -160,7 +166,12 @@ let rec compile_expr env : expr -> (environ * lean_expression) = function
     (match t with
      (* Call is a special case, being 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; assigned = true } in
+       (* Assumption: `name' is in the `env' because of our up-front variable analysis, but we've marked it as unassigned. *)
+       let { lean_id; complete; assigned } = Map.find_exn env name in
+       assert (not assigned) ;
+
+       let variable = { assigned = true; lean_id; complete } in
+       let env' = Map.set env ~key:name ~data:variable in
        (
          env',
          compile_term env' (Call ((Var "eq", []), [ Var name; rhs ]))
@@ -206,6 +217,8 @@ and compile_literal env : literal -> (environ * lean_expression) = function
 and compile_query env literals : (environ * lean_expression) =
   let { free; bound; _ } = extract_variables env literals in
   let env' = List.map free ~f:(fun name -> (name, { lean_id = name; complete = true; assigned = true }))
+             |> List.append
+               (List.map bound ~f:(fun name -> (name, { lean_id = name; complete = true; assigned = false })))
              |> Environ.of_alist_exn
              |> Map.merge env ~f:(fun ~key:_key -> function
                  | `Right v -> Some v