Compiles examples/01-computable-assign-scalar.rego

5 files changed, 141 insertions, 12 deletions

diff --git a/rego-proofs/compiler-rewrite/bin/main.ml b/rego-proofs/compiler-rewrite/bin/main.ml
index 619527e..933d97a 100644
--- a/rego-proofs/compiler-rewrite/bin/main.ml
+++ b/rego-proofs/compiler-rewrite/bin/main.ml
@@ -1,7 +1,13 @@
 open Core
 open Compiler_rewrite.Ast
+open Compiler_rewrite.Compiler
+open Compiler_rewrite.Lean
 
 let () =
   let file_name = Array.get (Sys.get_argv ()) 1 in
   let policy_code = In_channel.read_all file_name in
-  ast_of_string policy_code |> show_package |> print_endline
+  ast_of_string policy_code
+  |> compile_package
+  |> List.concat
+  |> List.map ~f:fmt_lean_expression
+  |> List.iter ~f:print_endline
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 01e076d..a0f521b 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -69,6 +69,8 @@ let term_of_ref_target : ref_target -> term = function
   | SetComprehension (key, q)    -> SetComprehension (key, q)
   | Call (fn, args)              -> Call (fn, args)
 
+let literal_of_expr exp : literal = Expr (exp, None)
+
 type compile_term_result =
   { form : lean_expression;
     toplevel : lean_expression list;
@@ -155,3 +157,117 @@ let rec compile_term env : term -> compile_term_result =
     { form = Application (form head, List.map ~f:form subterms);
       toplevel = toplevel head @ (List.map ~f:toplevel subterms |> List.concat);
       new_locals = new_locals head @ (List.map ~f:new_locals subterms |> List.concat) }
+
+let rec compile_expr env : expr -> compile_term_result = function
+  | Term t -> compile_term env t
+  | Every (Some key, value, domain, query) ->
+    (match compile_term env domain with
+     | { form = domain'; toplevel; new_locals } ->
+       let (toplevel', query_compiled) = compile_query env query in
+       let form = ForAll
+           (([key; value], "Rego.Any"),
+            (Application
+               (Identifier "And",
+                [(Application
+                    ((Identifier "Membership.mem"),
+                     [Tuple [Identifier key; Identifier value]; domain']));
+                 query_compiled])))
+       in {form; toplevel = toplevel @ toplevel'; new_locals})
+  | Every (None, value, domain, query) ->
+    (match compile_term env domain with
+     | { form = domain'; toplevel; new_locals } ->
+       let (toplevel', query_compiled) = compile_query env query in
+       let form = ForAll
+           (([value], "Rego.Any"),
+            (Application
+               (Identifier "And",
+                [(Application
+                    ((Identifier "Membership.mem"),
+                     [Identifier value; domain']));
+                 query_compiled])))
+       in {form; toplevel = toplevel @ toplevel'; new_locals})
+
+and compile_literal env : literal -> compile_term_result = function
+  | Exists (exprs, _) ->
+    let (toplevel, form) = compile_query env (List.map ~f:literal_of_expr exprs) in
+    { form; toplevel; new_locals = [] }
+  | Expr (expr, _) -> compile_expr env expr
+  | Negate (_expr, _) -> failwith "unimplemented"
+
+and compile_query env literals : (lean_expression list * lean_expression) =
+  let subterms = List.map ~f:(compile_literal env) literals in
+  let (compiled_terms, toplevel, new_locals) =
+    List.fold subterms
+      ~init:([], [], [])
+      ~f:(fun (a, b, c) y ->
+          match y with
+          | { form; toplevel; new_locals } ->
+            let comparison =
+              Application (Identifier "Ne", [form; Application (Identifier "Rego.Any.bool", [Bool false])])
+            in
+            (a @ [comparison], b @ toplevel, c @ new_locals))
+  in
+  (toplevel,
+   if List.is_empty new_locals then
+     match compiled_terms with
+     | [x] -> x
+     | xs -> Application (Identifier "And", xs)
+   else
+     Exists ((List.map ~f:fst new_locals, "Rego.Any"),
+             Application (Identifier "And",
+                          List.map ~f:(fun (x, y) ->
+                              Application (Identifier y, [Identifier x]))
+                            new_locals
+                          @ compiled_terms)))
+
+(* TODO: Probably get rid of this once we deal with nested definitions. *)
+let string_of_reference : reference -> string = function
+  | (Var name, _) -> name
+  | _ -> failwith "unimplemented"
+
+let compile_definition (rules : rule list) =
+  assert (not (List.is_empty rules));
+
+  let variable_being_defined =
+    match List.nth_exn rules 0 with
+    | { head = (rf, _); _ } -> rf
+  in
+
+  Inductive
+    (string_of_reference variable_being_defined,
+     [([gensym "x"], "Rego.Any")],
+     List.mapi rules ~f:(fun i ->
+         function
+         | { head = (_, Assignment (_, term)); body; _ } ->
+           (match compile_term Environ.empty term with
+            (* TODO: We need to pass this up.. not ignore it... *)
+            | { form; _ } ->
+              (Printf.sprintf "rule%d" i,
+               compile_query Environ.empty body |> snd,
+               [form]))
+         | _ -> failwith "unimplemented"))
+
+let compile_policy (rules : rule list) =
+  let group accessor l =
+    let rec grouping acc = function
+      | [] -> acc
+      | hd::tl ->
+        let l1,l2 = List.partition_tf ~f:(fun x -> equal_reference (accessor x) (accessor hd)) tl in
+        grouping ((hd::l1)::acc) l2
+    in
+    grouping [] l in
+  let head_of_rule = function | {head = (rf, _); _ } -> rf in
+  let by_reference = group head_of_rule rules in
+
+  (* let _scope = List.map by_reference *)
+  (*     ~f:(fun x -> *)
+  (*         let lean_id = List.nth_exn x 0 |> head_of_rule |> string_of_reference in *)
+  (*         lean_id *)
+  (*       ) *)
+  (* in *)
+
+  List.map by_reference ~f:compile_definition
+
+let compile_package = function
+  | { policy; _ } ->
+    [(* Comment path; Comment imports *)compile_policy policy]
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.mli b/rego-proofs/compiler-rewrite/lib/compiler.mli
index 193376c..66bd4dc 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.mli
+++ b/rego-proofs/compiler-rewrite/lib/compiler.mli
@@ -1,5 +1,7 @@
 open Ast
+open Lean
 
 type environ
 
 val extract_local_variables_wrapper : query -> string list
+val compile_package : package -> lean_expression list list
diff --git a/rego-proofs/compiler-rewrite/lib/lean.ml b/rego-proofs/compiler-rewrite/lib/lean.ml
index 975b563..f43c5c3 100644
--- a/rego-proofs/compiler-rewrite/lib/lean.ml
+++ b/rego-proofs/compiler-rewrite/lib/lean.ml
@@ -1,6 +1,7 @@
 open Core
 
 type lean_expression =
+  | Comment of string
   | Identifier of string
   | Bool of bool
   | String of string
@@ -9,7 +10,7 @@ type lean_expression =
   | Tuple of lean_expression list
   | Application of lean_expression * lean_expression list
   | Define of string * lean_expression
-  | Inductive of string * judgement list * (string * lean_expression) list
+  | Inductive of string * judgement list * (string * lean_expression * lean_expression list) list
   | ForAll of judgement * lean_expression
   | Exists of judgement * lean_expression
 
@@ -19,6 +20,7 @@ let fmt_judgement : judgement -> string = function
   | vars, t -> Printf.sprintf "(%s : %s)" (String.concat ~sep:" " vars) t
 
 let rec fmt_lean_expression : lean_expression -> string = function
+  | Comment s -> Printf.sprintf "\n-- %s\n" s
   | Identifier s -> s
   | Bool b -> if b then "True" else "False"
   | String s -> Printf.sprintf "\"%s\"" s
@@ -39,13 +41,15 @@ let rec fmt_lean_expression : lean_expression -> string = function
   | Exists (j, expr) ->
       Printf.sprintf "∃ %s, %s" (fmt_judgement j) (fmt_lean_expression expr)
   | Inductive (name, judgements, arms) ->
-      let fmt_arm (name, expr) =
-        let head =
-          Printf.sprintf "%s %s" name
-            (List.map ~f:fst judgements |> List.concat |> String.concat ~sep:" ")
-        in
-        Printf.sprintf "  | %s -> %s : %s" name (fmt_lean_expression expr) head
+    let fmt_arm (arm_name, expr, values) =
+      let head =
+        Printf.sprintf
+          "%s %s"
+          name
+          (List.map ~f:fmt_lean_expression values |> String.concat ~sep:" ")
       in
-      Printf.sprintf "inductive %s %s where\n%s" name
-        (List.map ~f:fmt_judgement judgements |> String.concat ~sep:" ")
-        (List.map ~f:fmt_arm arms |> String.concat ~sep:"\n")
+      Printf.sprintf "  | %s : %s → %s" arm_name (fmt_lean_expression expr) head
+    in
+    Printf.sprintf "inductive %s %s where\n%s" name
+      (List.map ~f:fmt_judgement judgements |> String.concat ~sep:" ")
+      (List.map ~f:fmt_arm arms |> String.concat ~sep:"\n")
diff --git a/rego-proofs/compiler-rewrite/lib/lean.mli b/rego-proofs/compiler-rewrite/lib/lean.mli
index 3b2b10c..1626dfa 100644
--- a/rego-proofs/compiler-rewrite/lib/lean.mli
+++ b/rego-proofs/compiler-rewrite/lib/lean.mli
@@ -1,4 +1,5 @@
 type lean_expression =
+  | Comment of string
   | Identifier of string
   | Bool of bool
   | String of string
@@ -7,7 +8,7 @@ type lean_expression =
   | Tuple of lean_expression list
   | Application of lean_expression * lean_expression list
   | Define of string * lean_expression
-  | Inductive of string * judgement list * (string * lean_expression) list
+  | Inductive of string * judgement list * (string * lean_expression * lean_expression list) list
   | ForAll of judgement * lean_expression
   | Exists of judgement * lean_expression