More reliable handling of rules with formal parameters

1 files changed, 136 insertions, 44 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 9b0e4ff..f1cc462 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -201,7 +201,9 @@ let rec compile_term env : term -> lean_expression =
   | Ref (target, terms) ->
     let target' = compile_term env (term_of_ref_target target) in
     let terms' = List.map ~f:(compile_term env) terms in
-    Application (Identifier "Rego.Any.deref", [target'] @ terms')
+    if List.is_empty terms' then
+      target'
+    else Application (Identifier "Rego.Any.deref", [target'] @ terms')
   | Array xs ->
     let xs' = List.map ~f:(compile_term env) xs in
     Application (Identifier "Rego.Any.array", xs')
@@ -227,6 +229,50 @@ let rec compile_term env : term -> lean_expression =
     let args' = List.map ~f:(compile_term env) args in
     Application (fn', args')
 
+let rewrite_with (env : environ) (q : query) : ((string * string) list * lean_expression list) =
+  let empty = ([], []) in
+  let combine (x1, y1) (x2, y2) = (x1 @ x2, y1 @ y2) in
+  let rec handle_term substitutions : term -> ((string * string) list * lean_expression list) = function
+    | Var name as v ->
+      (match List.Assoc.find ~equal:equal_term substitutions v with
+       | Some value ->
+         let new_identifier = gensym name in
+         ([(name, new_identifier)],
+          [Application (Identifier "Eq", [Identifier new_identifier;
+                                          compile_term env value])])
+       | None -> empty)
+    | Ref (target, args) ->
+      let target = term_of_ref_target target in
+      let terms = [target] @ args in
+      List.map terms ~f:(handle_term substitutions) |> List.fold ~init:empty ~f:combine
+    | Array xs | Set xs ->
+      List.map xs ~f:(handle_term substitutions) |> List.fold ~init:empty ~f:combine
+    | Object kvs -> List.map ~f:snd kvs
+                    |> List.map ~f:(handle_term substitutions)
+                    |> List.fold ~init:empty ~f:combine
+    | Call (x, xs) ->
+      List.map ((Ref x : term)::xs) ~f:(handle_term substitutions) |> List.fold ~init:empty ~f:combine
+    | ArrayComprehension _ -> assert false (* TODO *)
+    | ObjectComprehension _ -> assert false (* TODO *)
+    | SetComprehension _ -> assert false (* TODO *)
+    | _ -> empty
+  and handle_expr substitutions = function
+    | Term t -> handle_term substitutions t
+    | Every (_, _, _, q) ->
+      List.map q ~f:handle_literal |> List.fold ~init:empty ~f:combine
+  and handle_literal : literal -> ((string * string) list * lean_expression list) = function
+    | Exists (exprs, Some with_statement) ->
+      List.map exprs ~f:(handle_expr with_statement)
+        |> List.fold ~init:empty ~f:combine
+    | Expr (expr, Some with_statement) ->
+      handle_expr with_statement expr
+    | Negate (expr, Some with_statement) ->
+      handle_expr with_statement expr
+    | _ ->
+      empty
+  in
+  List.map ~f:handle_literal q |> List.fold ~init:empty ~f:combine
+
 let rec compile_expr env : expr -> (environ * lean_expression) = function
   | Term t ->
     (match t with
@@ -280,64 +326,121 @@ and compile_literal env : literal -> (environ * lean_expression) = function
   | Expr (expr, _) ->
     let (env', inner) = compile_expr env expr in
     (env', Application (Identifier "Ne", [inner; Application (Identifier "Rego.Any.bool", [Bool false])]))
-  | Negate (expr, _) ->
+   | Negate (expr, _) ->
     let (env', form) = compile_literal env (Expr (expr, None)) in
     (env', Application (Identifier "Not", [form]))
 
 and compile_query env literals : (environ * lean_expression) =
-  let { free; bound; used } = extract_variables env literals in
+  let (env', forms) = map_with_environment_updates env literals compile_literal in
+  let form = match forms with
+    | [x] -> x
+    | xs -> Application (Identifier "And", xs) in
+  (env', form)
+
+and compile_rule env rule : (environ * lean_expression) =
+  let { body = literals; _ } = rule in
+  let (new_identifiers, _with_statement_constraints) = rewrite_with env literals in
+  let env = List.map new_identifiers
+              ~f:(fun (old_name, anonymous_name) ->
+                (old_name, { lean_id = anonymous_name; kind = Parameter; assigned = true }))
+            |> extend_environment env in
+
+  let subjects =
+    let identifier_pair_of_term default_name : term -> _ = function
+      | Var identifier
+           (* This variable is only an identifier if it's actually used in the
+              body. This is a pretty awful way of checking if that's the case,
+              but it means we have a less complex usage analysis. *)
+           when
+             let env' = [(identifier, { lean_id = "unimportant"; kind = Parameter; assigned = true })]
+                        |> extend_environment env in
+             let { used; _ } = extract_variables env' literals in
+             List.mem used identifier ~equal:String.equal ->
+         let anon_name = gensym identifier in
+         (identifier, anon_name, Parameter, true, None)
+      | term ->
+         let anon_name = gensym default_name in
+         let assertion = (Call ((Var "eq", []), [Var anon_name; term]) : term) in
+         (anon_name, anon_name, Complete, true, Some assertion)
+    in
+    match rule |> kind_and_args_of_rule with
+    | `ScalarRule term | `SetRule term ->
+       [term |> identifier_pair_of_term "x"]
+    | `ObjectRule (key, value) ->
+       [key |> identifier_pair_of_term "k";
+        value |> identifier_pair_of_term "v"]
+    | `FunctionRule (args, output) ->
+       List.map args ~f:(fun arg_name -> (arg_name, gensym arg_name, Parameter, false, None))
+       @ [output |> identifier_pair_of_term "y"]
+  in
+  let env' =
+    List.map subjects
+      ~f:(fun (subject, lean_id, kind, output, _) ->
+        match kind with
+        | Parameter -> (subject, { lean_id = lean_id; assigned = not output; kind })
+        | _ -> (subject, { lean_id = lean_id; assigned = true; kind }))
+    |> extend_environment env
+  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
+                                            | { kind = Parameter; _ } -> false (* At what cost? *)
                                             | _ -> true) in
   let env' =
     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
+    |> extend_environment env' in
 
-  let lean_id_of_identifier k = (let { lean_id; _ } = Map.find_exn env'' k in lean_id) in
+  let lean_id_of_identifier k = (let { lean_id; _ } = Map.find_exn env' k in lean_id) in
   let constraints =
     List.filter_map used ~f:(fun name ->
         match Map.find env name with
         | Some { lean_id; kind = ScalarRule; _ } ->
-          Some
-            (Application (Identifier "scalar_specified_by",
-                          [Identifier lean_id;
-                           Identifier (lean_id_of_identifier name)]))
+           Some
+             (Application (Identifier "scalar_specified_by",
+                           [Identifier lean_id;
+                            Identifier (lean_id_of_identifier name)]))
         | Some { lean_id; kind = SetRule; _ } ->
-          Some
-            (Application (Identifier "set_specified_by",
-                          [Identifier lean_id;
-                           Identifier (lean_id_of_identifier name)]))
+           Some
+             (Application (Identifier "set_specified_by",
+                           [Identifier lean_id;
+                            Identifier (lean_id_of_identifier name)]))
         | Some { lean_id; kind = ObjectRule; _ } ->
-          Some
-            (Application (Identifier "object_specified_by",
-                          [Identifier lean_id;
-                           Identifier (lean_id_of_identifier name)]))
+           Some
+             (Application (Identifier "object_specified_by",
+                           [Identifier lean_id;
+                            Identifier (lean_id_of_identifier name)]))
         | Some { lean_id; kind = FunctionRule; _ } ->
-          Some
-            (Application (Identifier "function_specified_by",
-                          [Identifier lean_id;
-                           Identifier (lean_id_of_identifier name)]))
+           Some
+             (Application (Identifier "function_specified_by",
+                           [Identifier lean_id;
+                            Identifier (lean_id_of_identifier name)]))
         | Some _ -> None
         | None -> failwith (Printf.sprintf "unbound variable %s" name)) in
-  let compiled_terms = constraints @ subterms in
 
-  let form = match compiled_terms with
-     | [x] -> x
-     | xs -> Application (Identifier "And", xs)
-  in
 
+  let (env'', form) = compile_query env' literals in
   let quantified_names = free @ bound @ used in
+
+  let subject_constraints =
+    List.filter_map subjects ~f:(fun (_, _, _, _, x) -> x)
+    |> List.map ~f:(fun clause ->
+           let (_, compiled) = [Expr (Term clause, None)] |> compile_query env' in
+           compiled) in
+
+
   (env'',
    if List.is_empty quantified_names then
-     form
+     Application (Identifier "And", subject_constraints @ [form])
    else
      Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
-             form))
+             Application (Identifier "And", subject_constraints @ constraints @ [form])))
+
 
 let compile_definition env (rules : rule list) =
   assert (not (List.is_empty rules));
@@ -394,25 +497,14 @@ let compile_definition env (rules : rule list) =
   in
   let conditional_arms =
     List.mapi conditional_rules ~f:(fun i rule ->
-        let { body; _ } = rule in
-        let (_, body_form) = compile_query env' body in
-        let subject_constraints =
-          List.filter_map subjects ~f:(fun (_, _, _, _, x) -> x)
-          |> List.map ~f:(fun clause ->
-              let (_, compiled) = [Expr (Term clause, None)] |> compile_query env' in
-              compiled) in
-        (body_form,
+        let (_env', body) = compile_rule env' rule in
+        (body,
          (Printf.sprintf "rule%d" i,
-          subject_constraints @ [body_form],
+          [body],
           List.map subjects ~f:(fun (_, x, _, _, _) -> Identifier x)))) in
   let default_arms =
     List.map default_rules ~f:(fun _ ->
         let other_rule_bodies = List.map ~f:fst conditional_arms in
-        let subject_constraints =
-          List.filter_map subjects ~f:(fun (_, _, _, _, x) -> x)
-          |> List.map ~f:(fun clause ->
-              let (_, compiled) = [Expr (Term clause, None)] |> compile_query env' in
-              compiled) in
         let body_form =
           let fold_or props =
             match List.reduce props ~f:(fun x y -> Application (Identifier "Or", [y; x])) with
@@ -422,7 +514,7 @@ let compile_definition env (rules : rule list) =
           Application (Identifier "Not", [fold_or other_rule_bodies])
         in
         ("rule_default",
-         subject_constraints @ [body_form],
+         [body_form],
          List.map subjects ~f:(fun (_, x, _, _, _) -> Identifier x))) in
 
   let form =