Handle nested parameters

2 files changed, 20 insertions, 22 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 8cf3cb6..06d4040 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -400,18 +400,21 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
         subjects
         candidates
         ~f:(fun subject ->
-          function
-          | Var identifier when List.mem bound identifier ~equal:String.equal ->
-             (* In this case, we're dealing with a formal parameter, so we want
+            function
+            (* In this case, we're dealing with a formal parameter, so we want
                 to have it be an alias of the lean_id for the subject. *)
-             (Some identifier,
-              (identifier, { lean_id = subject; kind = Parameter; assigned = true }),
-              None)
-          | term ->
-             let anon_name = gensym "x" in
-             (None,
-              (anon_name, { lean_id = anon_name; kind = Parameter; assigned = true }),
-              Some (Call ((Var "eq", []), [Var anon_name; term]) : term)))
+            | Var identifier when List.mem bound identifier ~equal:String.equal ->
+              (Some identifier,
+               [(identifier, { lean_id = subject; kind = Parameter; assigned = true })],
+               None)
+            | term ->
+              let { free; _ } = extract_variables env [literal_of_term term] in
+              (None,
+               [(subject, { lean_id = subject; kind = Parameter; assigned = true })]
+               @ List.map free
+                 ~f:(fun name ->
+                     (name, { lean_id = gensym "param"; kind = Parameter; assigned = true })),
+               Some (Call ((Var "eq", []), [Var subject; term]) : term)))
     in
 
     let parameters = env_updates_and_constraints |> List.filter_map ~f:(fun (x, _, _) -> x) in
@@ -424,7 +427,7 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
 
     (non_parameters,
      parameters,
-     env_updates_and_constraints |> List.map ~f:(fun (_, x, _) -> x),
+     env_updates_and_constraints |> List.map ~f:(fun (_, x, _) -> x) |> List.concat,
      env_updates_and_constraints |> List.filter_map  ~f:(fun (_, _, x) -> x))
   in
   let env' = extend_environment env env_updates in
@@ -469,22 +472,17 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
         | None -> failwith (Printf.sprintf "unbound variable %s" name)) in
 
 
-  let (env'', form) = compile_query env' literals in
+  let subject_constraints = List.map subject_constraints ~f:literal_of_term in
+  let (env'', form) = compile_query env' (literals @ subject_constraints) in
   let quantified_names = free @ bound @ used in
 
-  let subject_constraints =
-    List.map subject_constraints
-      ~f:(fun clause ->
-        let (_, compiled) = [literal_of_term clause] |> compile_query env' in
-        compiled) in
-
   (env'',
    if List.is_empty quantified_names then
-     Application (Identifier "And", subject_constraints @ [form])
+     Application (Identifier "And", [form])
      |> simplify_expression
    else
      Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
-             Application (Identifier "And", subject_constraints @ constraints @ [form])
+             Application (Identifier "And", constraints @ [form])
              |> simplify_expression))
   :: (match rule.otherwise with
       | Some branch -> compile_rule env subjects branch
diff --git a/rego-proofs/examples/13-else-keyword.rego b/rego-proofs/examples/13-else-keyword.rego
index f04887a..f8a4b9d 100644
--- a/rego-proofs/examples/13-else-keyword.rego
+++ b/rego-proofs/examples/13-else-keyword.rego
@@ -2,4 +2,4 @@ package examples.set_defined_by_predicate
 
 import future.keywords.if
 
-phi := 3.14159 if False else := 1.61803
+phi := 3.14159 if false else := 1.61803