1 files changed, 134 insertions, 0 deletions
diff --git a/rego-proofs/compiler/test/compiler.ml b/rego-proofs/compiler/test/compiler.ml
new file mode 100644
index 0000000..741ae3a
--- /dev/null
+++ b/rego-proofs/compiler/test/compiler.ml
@@ -0,0 +1,134 @@
+open Core
+open Compiler
+open Compiler__
+
+module To_test = struct
+  let extract_variables = Compiler.extract_variables_wrapper
+  let transform_ref_head = Compiler.transform_rule_with_ref_head
+  let fmt = Lean.fmt_lean_expression
+end
+
+let strings =
+  Alcotest.testable
+    (Format.pp_print_list Format.pp_print_string)
+    (List.equal String.equal)
+
+let extract_query policy =
+  match Ast.ast_of_string policy with
+  | { policy = [ { body; _ } ]; _ } -> body
+  | _ -> failwith "failed to parse"
+
+let test_local_extraction_simple_1 () =
+  let policy_code = {|
+package test
+
+q {
+  x := 2
+}
+  |} in
+  let { bound; _ } = To_test.extract_variables (extract_query policy_code) in
+  Alcotest.(check strings)
+    "expected identifier list" [ "x" ] bound
+
+let test_local_extraction_simple_2 () =
+  let policy_code = {|
+package test
+
+q {
+  x := 2
+  y := 3
+  z := 4
+}
+  |} in
+  let { bound; _ } = To_test.extract_variables (extract_query policy_code) in
+  Alcotest.(check strings)
+    "expected identifier list" [ "x"; "y"; "z" ] bound
+
+let test_local_dependent_on_previously_defined_local () =
+  let policy_code = {|
+package test
+
+q {
+  x := 2
+  y := x
+}
+  |} in
+  let { bound; _ } = To_test.extract_variables (extract_query policy_code) in
+  Alcotest.(check strings)
+    "expected identifier list" [ "x"; "y" ] bound
+
+let test_free_extraction_simple_1 () =
+  let policy_code = {|
+package test
+
+q {
+  y := 7
+  y == 8
+}
+  |} in
+  let { free; bound; _ } = To_test.extract_variables (extract_query policy_code) in
+  Alcotest.(check strings)
+    "expected identifier list" [] free;
+  Alcotest.(check strings)
+    "expected identifier list" [ "y" ] bound
+
+let test_free_extraction_simple_2 () =
+  let policy_code = {|
+package test
+
+q {
+  y := z
+  z = 7
+}
+  |} in
+  let { free; bound; _ } = To_test.extract_variables (extract_query policy_code) in
+  Alcotest.(check strings)
+    "expected identifier list" [ "z" ] free;
+  Alcotest.(check strings)
+    "expected identifier list" [ "y" ] bound
+
+let test_identifier () =
+  let open Lean in
+  let ast = Identifier "test" in
+  Alcotest.(check string) "expected string" "test" (To_test.fmt ast)
+
+let test_transform_simple_1 () =
+  let open Ast in
+  let rule = { head = ((Var "fruit", [Var "apple"; Var "seeds"]), Assignment (Assign, (Scalar (Number 12.))));
+               default = false;
+               body = [Expr (Term (Scalar (Bool true)), None)];
+               otherwise = None } in
+  Alcotest.(check int) "expected rules"
+    2
+    (To_test.transform_ref_head rule |> List.length)
+
+let () =
+  let open Alcotest in
+  run "Compiler"
+    [
+      ( "free variable extraction",
+        [
+          test_case "free variable extraction simple (1)" `Quick
+            test_free_extraction_simple_1;
+          test_case "free variable extraction simple (2)" `Quick
+            test_free_extraction_simple_2;
+        ] );
+      ( "local extraction",
+        [
+          test_case "local extraction simple (1)" `Quick
+            test_local_extraction_simple_1;
+          test_case "local extraction simple (2)" `Quick
+            test_local_extraction_simple_2;
+          test_case
+            "local extraction with local dependent on previously-defined local \
+             (1)"
+            `Quick test_local_dependent_on_previously_defined_local
+        ] );
+      ( "nested definitions",
+        [
+          test_case "local nested definitions (1)" `Quick
+            test_transform_simple_1;
+        ] ) ;
+      ( "ast formatting ",
+        [ test_case "identifier compilation" `Quick test_identifier ] );
+    ]