1 files changed, 46 insertions, 14 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 1cb325e..9044f50 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -240,24 +240,53 @@ and compile_literal env : literal -> (environ * lean_expression) = function
   | Negate (_expr, _) -> failwith "unimplemented"
 
 and compile_query env literals : (environ * lean_expression) =
-  let { free; bound; _ } = extract_variables env literals in
+  let { free; bound; used } = extract_variables env literals in
+  let used = used
+             |> List.filter ~f:(fun name -> match Map.find_exn env name with
+                                            | { kind = Complete; _ } -> false
+                                            | _ -> true) in
   let env' =
-    List.append
-      (List.map free  ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = true })))
-      (List.map bound ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = false })))
+    List.concat
+      [List.map free  ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = true }));
+       List.map used  ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = true }));
+       List.map bound ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = false }))]
     |> extend_environment env in
   let (env'', subterms) = map_with_environment_updates env' literals compile_literal in
   let compiled_terms = List.map subterms
       ~f:(fun x -> Application (Identifier "Ne", [x; Application (Identifier "Rego.Any.bool", [Bool false])])) in
-  let quantified_names = free @ bound in
+
+  let lean_id_of_identifier k = (let { lean_id; _ } = Map.find_exn env'' k in lean_id) in
+  let constraints =
+    List.map bound ~f:(fun name ->
+        match Map.find env name with
+        | Some { lean_id; kind = ScalarRule; _ } ->
+           Application (Identifier "scalar_specified_by",
+                        [Identifier lean_id;
+                         Identifier (lean_id_of_identifier name)])
+        | Some { lean_id; kind = SetRule; _ } ->
+           Application (Identifier "set_specified_by",
+                        [Identifier lean_id;
+                         Identifier (lean_id_of_identifier name)])
+        | Some { lean_id; kind = ObjectRule; _ } ->
+           Application (Identifier "object_specified_by",
+                        [Identifier lean_id;
+                         Identifier (lean_id_of_identifier name)])
+        | Some { lean_id; kind = FunctionRule; _ } ->
+           Application (Identifier "function_specified_by",
+                        [Identifier lean_id;
+                         Identifier (lean_id_of_identifier name)])
+        | None -> failwith (Printf.sprintf "unbound variable %s" name)
+        | _ -> failwith "unreachable") in
+  let compiled_terms = constraints @ compiled_terms in
+
+  let quantified_names = free @ bound @ used in
   (env'',
    if List.is_empty quantified_names then
      match compiled_terms with
      | [x] -> x
      | xs -> Application (Identifier "And", xs)
    else
-     let lean_id_of_identifier k = (let { lean_id; _ } = Map.find_exn env'' k in lean_id) in
-     Exists ((List.map (free @ bound) ~f:lean_id_of_identifier , "Rego.Any"),
+     Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
              Application (Identifier "And", compiled_terms)))
 
 (* TODO: Probably get rid of this once we deal with nested definitions. *)
@@ -273,17 +302,21 @@ let compile_definition env (rules : rule list) =
     | { head = (rf, _); _ } -> rf
   in
 
+  let subject = gensym "x" in
   let form =
     Inductive
       (string_of_reference variable_being_defined,
-       [([gensym "x"], "Rego.Any")],
+       [([subject], "Rego.Any")],
        List.mapi rules ~f:(fun i ->
            function
            | { head = (_, Assignment (_, term)); body; _ } ->
-             let form = compile_term env term in
-             (Printf.sprintf "rule%d" i,
-              compile_query env body |> snd,
-              [form])
+              let eq : term = Call ((Var "eq", []), [Var subject; term]) in
+              let env' = [(subject, { lean_id = subject; kind = Complete; assigned = true })]
+                         |> extend_environment env in
+              let (_, form) = compile_query env' [Expr (Term eq, None)] in
+              (Printf.sprintf "rule%d" i,
+               [form; compile_query env body |> snd],
+               [Identifier subject])
            | _ -> failwith "unimplemented"))
   in (env, form)
 
@@ -312,7 +345,6 @@ let compile_package = function
   | { path; imports; policy } ->
     let policy = List.map ~f:transform_rule_with_ref_head policy |> List.concat in
 
-    let env = toplevel in
     let env =
       List.map policy ~f:(function
           | { head = ((Var name, []), Assignment _); _ } ->
@@ -324,7 +356,7 @@ let compile_package = function
           | { head = ((Var name, []), Function _); _ } ->
              (name, { lean_id = name ; kind = FunctionRule; assigned = true })
           | _ -> failwith "policy rule head is not a variable")
-      |> extend_environment env in
+      |> extend_environment toplevel in
 
     let (_env', forms) = compile_policy env policy in
     [Comment (show_reference path |> String.escaped);