Compile error condition

3 files changed, 242 insertions, 49 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 06d4040..898ff51 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -11,6 +11,16 @@ type definition_kind = Complete
 [@@deriving eq, show]
 
 
+let fold_or props =
+  match List.reduce props ~f:(fun x y -> Application (Identifier "Or", [y; x])) with
+  | Some x -> x
+  | None -> Bool true
+
+let fold_and props =
+  match List.reduce props ~f:(fun x y -> Application (Identifier "And", [y; x])) with
+  | Some x -> x
+  | None -> Bool false
+
 let kind_and_args_of_rule = function
   | { head = ((_, []), Assignment (_, term)); _ } ->
     `ScalarRule term
@@ -75,7 +85,20 @@ let term_of_ref_target : ref_target -> term = function
   | SetComprehension (key, q)    -> SetComprehension (key, q)
   | Call (fn, args)              -> Call (fn, args)
 
+let ref_target_of_term : term -> ref_target = function
+  | Var s                        -> Var s
+  | Array xs                     -> Array xs
+  | Object kvs                   -> Object kvs
+  | Set xs                       -> Set xs
+  | ArrayComprehension (key, q)  -> ArrayComprehension (key, q)
+  | ObjectComprehension (key, q) -> ObjectComprehension (key, q)
+  | SetComprehension (key, q)    -> SetComprehension (key, q)
+  | Call (fn, args)              -> Call (fn, args)
+  | _                            -> failwith "cannot widen to ref_target"
+
 let literal_of_expr exp : literal = Expr (exp, None)
+let literal_of_term t = Expr (Term t, None)
+
 
 let map_with_environment_updates env xs fn =
   List.fold xs ~init:(env, []) ~f:(fun (env, ys) x ->
@@ -134,9 +157,6 @@ let extract_variables (env : environ) (q : query) : analysis_result =
                     |> List.fold ~init:empty ~f:combine
     | Call (x, xs) ->
       List.map ((Ref x : term)::xs) ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
-    | ArrayComprehension _ -> assert false (* TODO *)
-    | ObjectComprehension _ -> assert false (* TODO *)
-    | SetComprehension _ -> assert false (* TODO *)
     | _ -> empty
   and handle_expr env = function
     | Term t -> handle_term env t
@@ -255,16 +275,103 @@ let rec compile_term env : term -> lean_expression =
      let kvs' = List.map ~f:compile_kv kvs in
      Application (Identifier "Rego.Any.object", kvs')
   | ArrayComprehension (_term, _query) ->
-     assert false (* TODO *)
+     failwith "ArrayComprehension encountered, should have been lifted"
   | SetComprehension (_term, _query) ->
-     assert false (* TODO *)
+     failwith "ArrayComprehension encountered, should have been lifted"
   | ObjectComprehension ((_key, _value), _query) ->
-     assert false (* TODO *)
+     failwith "ArrayComprehension encountered, should have been lifted"
   | Call (fn, args) ->
     let fn' = compile_term env (Ref fn) in
     let args' = List.map ~f:(compile_term env) args in
     Application (fn', args')
 
+let lift_term (env : environ) (rule : rule) (f : environ -> term -> ('a * term) option) : (('a * term) list * rule) =
+  let empty = ([], []) in
+  let combine = fun (side_info, q) (side_info', term) ->
+    (side_info @ side_info', q @ [term]) in
+  let rec handle_term : term -> (('a * term) list * term) = function
+    | term when f env term |> Option.is_some ->
+      let (side_info, term') = f env term |> Option.value_exn in
+      ([(side_info, term)], term')
+    | Ref (target, args) ->
+      let target = term_of_ref_target target in
+      let terms = [target] @ args in
+
+      let (side_info, terms') =
+        List.map terms ~f:handle_term
+        |> List.fold ~f:combine ~init:empty
+      in (side_info, Ref (ref_target_of_term (List.hd_exn terms'), List.tl_exn terms'))
+    | Array xs ->
+      let (side_info, terms) =
+        List.map xs ~f:handle_term
+        |> List.fold ~f:combine ~init:empty
+      in (side_info, Array terms)
+    | Set xs ->
+      let (side_info, terms) =
+        List.map xs ~f:handle_term
+        |> List.fold ~f:combine ~init:empty
+      in (side_info, Set terms)
+    | Object kvs ->
+      let (side_info, terms) =
+        List.map ~f:snd kvs
+        |> List.map  ~f:handle_term
+        |> List.fold ~f:combine ~init:empty
+      in
+      (side_info,
+       Object
+         (List.zip_exn
+            (List.map ~f:fst kvs)
+            terms))
+    | Call (x, xs) ->
+      let (side_info, terms) =
+        List.map ((Ref x : term)::xs) ~f:handle_term
+        |> List.fold ~f:combine ~init:empty
+      in
+      (match List.hd_exn terms with
+       | Ref reference ->
+         (side_info, Call (reference, List.tl_exn terms))
+       | term -> failwith (show_term term |> Printf.sprintf "invalid call target %s"))
+    | term -> ([], term)
+  and handle_expr = function
+    | Term t ->
+      let (side_info, t') = handle_term t in (side_info, Term t')
+    | Every (k, v, t, q) ->
+      let (side_info, q') =
+        List.map q ~f:handle_literal
+        |> List.fold ~f:(fun (side_info, q) (side_info', term) ->
+            (side_info @ side_info', q @ [term]))
+          ~init:empty
+      in (side_info, Every (k, v, t, q'))
+  and handle_literal : literal -> (('a * term) list * literal) = function
+    | Exists (exprs, with_statement) ->
+      let (side_info, exprs') =
+        List.map exprs ~f:handle_expr
+        |> List.fold ~init:empty ~f:combine
+      in (side_info, Exists (exprs', with_statement))
+    | Expr (expr, with_statement) ->
+      let (side_info, expr') = handle_expr expr in
+      (side_info, Expr (expr', with_statement))
+    | Negate (expr, with_statement) ->
+      let (side_info, expr') = handle_expr expr in
+      (side_info, Negate (expr', with_statement))
+  in
+  let { body = q; head; _ } = rule in
+  
+  let (side_info, query) = List.map ~f:handle_literal q |> List.fold ~init:empty ~f:combine in
+  let (side_info', head') = match head with
+    | (rf, Assignment (op, term)) ->
+      let (side_info, term') = handle_term term in
+      (side_info, (rf, Assignment (op, term')))
+    | (rf, Function (op, args, term)) ->
+      let (side_info, term') = handle_term term in
+      (side_info, (rf, Function (op, args, term')))
+    | (rf, Set term) ->
+      let (side_info, term') = handle_term term in
+      (side_info, (rf, Set term'))
+      in
+  (side_info @ side_info', { rule with body = query; head = head' })
+  
+
 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
@@ -340,7 +447,7 @@ let rec compile_expr env : expr -> (environ * lean_expression) = function
                 ((Identifier "Membership.mem"),
                  [Tuple [Identifier key; Identifier value]; domain']));
              query']))
-       |> simplify_expression))
+        |> simplify_expression))
   | Every (None, value, domain, query) ->
     let domain' = compile_term env domain in
     (* Extend the environment with the value. *)
@@ -355,7 +462,7 @@ let rec compile_expr env : expr -> (environ * lean_expression) = function
                 ((Identifier "Membership.mem"),
                  [Identifier value; domain']));
              query']))
-       |> simplify_expression))
+        |> simplify_expression))
 
 and compile_literal env : literal -> (environ * lean_expression) = function
   | Exists (exprs, _) ->
@@ -364,7 +471,7 @@ 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]))
 
@@ -375,24 +482,82 @@ and compile_query env literals : (environ * lean_expression) =
     | xs -> Application (Identifier "And", xs) |> simplify_expression in
   (env', form)
 
+and compile_comprehension env anonymous_name (comp : term) =
+  match comp with
+  | ArrayComprehension (value, q) ->
+    let { free; _ } = extract_variables env [literal_of_term value] in
+    let env = List.map free ~f:(fun identifier ->
+        (identifier, { lean_id = (gensym identifier); kind = Parameter; assigned = true }))
+              |> extend_environment env in
+    let (_, form) = compile_query env q in
+    Application (Identifier "array_specified_by", [Identifier anonymous_name; form])
+  | SetComprehension (value, q) ->
+    let { free; _ } = extract_variables env [literal_of_term value] in
+    let env = List.map free ~f:(fun identifier ->
+        (identifier, { lean_id = (gensym identifier); kind = Parameter; assigned = true }))
+              |> extend_environment env in
+    let (_, form) = compile_query env q in
+    Application (Identifier "set_specified_by", [Identifier anonymous_name; form])
+  | ObjectComprehension ((key, value), q) ->
+    let { free; _ } = extract_variables env [literal_of_term (term_of_object_key key)] in
+    let { free = free'; _ } = extract_variables env [literal_of_term value] in
+    let free = free @ free' in
+    let env = List.map free ~f:(fun identifier ->
+        (identifier, { lean_id = (gensym identifier); kind = Parameter; assigned = true }))
+              |> extend_environment env in
+    let (_, form) = compile_query env q in
+    Application (Identifier "object_specified_by", [Identifier anonymous_name; form])
+  | _ -> failwith (show_term comp |>
+                   (Printf.sprintf "called compile_comprehension on non-comprehension %s"))
+
 and compile_rule env subjects rule : (environ * lean_expression) list =
+  let (comprehensions, rule) = lift_term env rule
+      (fun _env term ->
+         let anonymous_name = gensym "comprehension" in
+         match term with
+         | ArrayComprehension _ | SetComprehension _ | ObjectComprehension _ ->
+           Some (anonymous_name, Var anonymous_name)
+         | _ -> None) in
+  (* Each comprehension is now mapped to an anonymous name. We'll add it to the
+     scope so we don't have an unbound variable error when we compile the
+     modified rule. *)
+  let env = List.map comprehensions
+      ~f:(fun (anonymous_name, _) ->
+          (anonymous_name, { lean_id = anonymous_name; kind = Parameter; assigned = true }))
+            |> extend_environment env in
+
   let { body = literals; _ } = rule in
+
+  (* TODO: Deal with user-defined functions. *)
+  (* let (_, literals) = lift_term env literals *)
+  (*     (fun env term -> *)
+  (*        match term with *)
+  (*        | Call ((Var identifier, []), _args) -> *)
+  (*          (match Map.find env identifier with *)
+  (*           | Some { kind = FunctionRule ; _ } -> *)
+  (*             let anon_name = gensym "function_call" in *)
+  (*             Some (anon_name, Var anon_name) *)
+  (*           | _ -> None) *)
+  (*        | _ -> None) in *)
+
+
+  List.iter literals ~f:(fun lit -> show_literal lit |> print_endline);
+
   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 }))
+      ~f:(fun (old_name, anonymous_name) ->
+          (old_name, { lean_id = anonymous_name; kind = Parameter; assigned = true }))
             |> extend_environment env in
 
-  let literal_of_term t = Expr (Term t, None) in
   let { free; bound; used } = extract_variables env literals in
   let (used', parameters, env_updates, subject_constraints) =
     let candidates = match rule |> kind_and_args_of_rule with
-    | `ScalarRule term | `SetRule term ->
-       [term]
-    | `ObjectRule (key, value) ->
-       [key; value]
-    | `FunctionRule (args, output) ->
-       List.map args ~f:(fun arg_name -> (Var arg_name : term)) @ [output]
+      | `ScalarRule term | `SetRule term ->
+        [term]
+      | `ObjectRule (key, value) ->
+        [key; value]
+      | `FunctionRule (args, output) ->
+        List.map args ~f:(fun arg_name -> (Var arg_name : term)) @ [output]
     in
 
     let env_updates_and_constraints =
@@ -419,11 +584,11 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
 
     let parameters = env_updates_and_constraints |> List.filter_map ~f:(fun (x, _, _) -> x) in
     let non_parameters = List.filter_map candidates ~f:(function
-                             | Var id ->
-                                if List.mem parameters id ~equal:String.equal then
-                                  None
-                                else Some id
-                             | _ -> None) in
+        | Var id ->
+          if List.mem parameters id ~equal:String.equal then
+            None
+          else Some id
+        | _ -> None) in
 
     (non_parameters,
      parameters,
@@ -435,8 +600,8 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
   let used = used @ used' in
   let used = used
              |> List.filter ~f:(fun name -> match Map.find_exn env name with
-                                            | { kind = Complete; _ } -> false
-                                            | _ -> true) in
+                 | { kind = Complete; _ } -> false
+                 | _ -> true) in
   let env' =
     List.concat
       [List.map free  ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = true }));
@@ -449,31 +614,32 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
     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 subject_constraints = List.map subject_constraints ~f:literal_of_term in
   let (env'', form) = compile_query env' (literals @ subject_constraints) in
+  let comprehension_constraints =
+    List.map comprehensions ~f:(fun (name, form) -> compile_comprehension env name form) in
   let quantified_names = free @ bound @ used in
 
   (env'',
@@ -482,7 +648,7 @@ and compile_rule env subjects rule : (environ * lean_expression) list =
      |> simplify_expression
    else
      Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
-             Application (Identifier "And", constraints @ [form])
+             Application (Identifier "And", constraints @ comprehension_constraints @ [form])
              |> simplify_expression))
   :: (match rule.otherwise with
       | Some branch -> compile_rule env subjects branch
@@ -535,11 +701,6 @@ let compile_definition env (rules : rule list) =
     List.map default_rules ~f:(fun _ ->
         let other_rule_bodies = List.map ~f:fst conditional_arms in
         let body_form =
-          let fold_or props =
-            match List.reduce props ~f:(fun x y -> Application (Identifier "Or", [y; x])) with
-            | Some x -> x
-            | None -> Bool true
-          in
           Application (Identifier "Not", [fold_or other_rule_bodies])
         in
         ("rule_default",
@@ -583,6 +744,32 @@ let compile_policy env (rules : rule list) =
     |> extend_environment env in
   map_with_environment_updates env' by_reference compile_definition
 
+let compile_error_condition env =
+  Map.to_alist env
+  |> List.filter_map ~f:(fun (k, v) ->
+      match v with
+      | { kind = ScalarRule; _ } ->
+        let anon_name = gensym "the" in
+        Some(
+          ExistsUnique (
+            ([anon_name], "Rego.Any"),
+            Application (Identifier "scalar_specified_by",
+                         [Identifier anon_name; Identifier k])))
+      | { kind = FunctionRule; _ } ->
+        let args_name = gensym "args" in
+        let value_name = gensym "value" in
+        Some(
+          ForAll
+            (([args_name], "Rego.Any"),
+             ExistsUnique (
+               ([value_name], "Rego.Any"),
+               Application
+                 (Identifier "eq",
+                  [Application (Identifier k, [Identifier args_name]);
+                   Identifier value_name]))))
+      | _ -> None)
+  |> fold_and
+  
 let compile_package =
     let string_of_reference : reference -> string = function
       | (Var name, xs) ->
@@ -608,9 +795,11 @@ let compile_package =
   function
   | { path; imports; policy } ->
     let policy = List.map ~f:transform_rule_with_ref_head policy |> List.concat in
-    let (_env', forms) = compile_policy toplevel policy in
+    let (env', forms) = compile_policy toplevel policy in
+    let error_form = Define ("rego_error", compile_error_condition env') in
     [Comment (Printf.sprintf "package %s"
                 (string_of_reference path)
               |> String.escaped)]
     @ List.map imports ~f:(fun x -> Comment (compile_import x))
     @ forms
+    @ [error_form]
diff --git a/rego-proofs/compiler-rewrite/lib/lean.ml b/rego-proofs/compiler-rewrite/lib/lean.ml
index b3526e3..6bb4265 100644
--- a/rego-proofs/compiler-rewrite/lib/lean.ml
+++ b/rego-proofs/compiler-rewrite/lib/lean.ml
@@ -13,6 +13,7 @@ type lean_expression =
   | Inductive of string * judgement list * (string * lean_expression list * lean_expression list) list
   | ForAll of judgement * lean_expression
   | Exists of judgement * lean_expression
+  | ExistsUnique of judgement * lean_expression
 [@@deriving eq, show]
 
 and judgement = string list * string [@@deriving eq, show]
@@ -41,6 +42,8 @@ let rec fmt_lean_expression : lean_expression -> string = function
       Printf.sprintf "∀ %s, %s" (fmt_judgement j) (fmt_lean_expression expr)
   | Exists (j, expr) ->
       Printf.sprintf "∃ %s, %s" (fmt_judgement j) (fmt_lean_expression expr)
+  | ExistsUnique (j, expr) ->
+      Printf.sprintf "∃! %s, %s" (fmt_judgement j) (fmt_lean_expression expr)
   | Inductive (name, judgements, arms) ->
     let fmt_arm (arm_name, exprs, values) =
       let head =
diff --git a/rego-proofs/compiler-rewrite/lib/lean.mli b/rego-proofs/compiler-rewrite/lib/lean.mli
index fb58c72..f6a769a 100644
--- a/rego-proofs/compiler-rewrite/lib/lean.mli
+++ b/rego-proofs/compiler-rewrite/lib/lean.mli
@@ -11,6 +11,7 @@ type lean_expression =
   | Inductive of string * judgement list * (string * lean_expression list * lean_expression list) list
   | ForAll of judgement * lean_expression
   | Exists of judgement * lean_expression
+  | ExistsUnique of judgement * lean_expression
 
 and judgement = string list * string