1 files changed, 27 insertions, 16 deletions
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index af0d14b..4756304 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -350,20 +350,7 @@ and compile_rule env subjects rule : (environ * lean_expression) =
             |> extend_environment env in
 
   let literal_of_term t = Expr (Term t, None) in
-  (* let non_query_forms = 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] *)
-  (* in *)
-  (* let forms = (literals @ (List.map ~f:literal_of_term non_query_forms)) 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 ->
@@ -558,10 +545,34 @@ let compile_policy env (rules : rule list) =
     |> extend_environment env in
   map_with_environment_updates env' by_reference compile_definition
 
-let compile_package = function
+let compile_package =
+    let string_of_reference : reference -> string = function
+      | (Var name, xs) ->
+        [name]
+        @ List.map xs ~f:(function
+            | Scalar (String x) -> x
+            | _ -> failwith "unreachable")
+      |> String.concat ~sep:"."
+
+      | _ -> assert false (* unreachable *)
+    in
+  let compile_import  = function
+    | { alias = None; path } ->
+      Printf.sprintf "imports %s"
+        (string_of_reference path)
+      |> String.escaped
+    | { alias = Some alias; path } ->
+      Printf.sprintf "imports %s as %s"
+        (string_of_reference path)
+        alias
+      |> String.escaped
+  in
+  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
-    [Comment (show_reference path |> String.escaped);
-     Comment (List.map ~f:show_import imports |> String.concat |> String.escaped)]
+    [Comment (Printf.sprintf "package %s"
+                (string_of_reference path)
+              |> String.escaped)]
+    @ List.map imports ~f:(fun x -> Comment (compile_import x))
     @ forms