1 files changed, 37 insertions, 17 deletions

diff --git a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
index 3de5d7e..ef42340 100644
--- a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
+++ b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
@@ -2,7 +2,7 @@ open Core
 open Compiler_rewrite
 
 module To_test = struct
-  let extract_locals = Compiler.extract_local_variables_wrapper
+  let extract_variables = Compiler.extract_variables_wrapper
   let fmt = Lean.fmt_lean_expression
 end
 
@@ -24,9 +24,9 @@ q {
   x := 2
 }
   |} in
+  let (_, bound) = To_test.extract_variables (extract_query policy_code) in
   Alcotest.(check strings)
-    "expected identifier list" [ "x" ]
-    (To_test.extract_locals (extract_query policy_code))
+    "expected identifier list" [ "x" ] bound
 
 let test_local_extraction_simple_2 () =
   let policy_code = {|
@@ -38,9 +38,9 @@ q {
   z := 4
 }
   |} in
+  let (_, bound) = To_test.extract_variables (extract_query policy_code) in
   Alcotest.(check strings)
-    "expected identifier list" [ "x"; "y"; "z" ]
-    (To_test.extract_locals (extract_query policy_code))
+    "expected identifier list" [ "x"; "y"; "z" ] bound
 
 let test_local_dependent_on_previously_defined_local () =
   let policy_code = {|
@@ -51,22 +51,39 @@ q {
   y := x
 }
   |} in
+  let (_, bound) = To_test.extract_variables (extract_query policy_code) in
   Alcotest.(check strings)
-    "expected identifier list" [ "x"; "y" ]
-    (To_test.extract_locals (extract_query policy_code))
+    "expected identifier list" [ "x"; "y" ] bound
 
-let test_local_not_ground_because_of_ordering () =
+let test_free_extraction_simple_1 () =
   let policy_code = {|
 package test
 
 q {
-  y := x
-  x := 2
+  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" [ "x" ]
-    (To_test.extract_locals (extract_query policy_code))
+    "expected identifier list" [ "y" ] bound
 
 let test_identifier () =
   let open Compiler_rewrite.Lean in
@@ -77,6 +94,13 @@ 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
@@ -86,11 +110,7 @@ let () =
           test_case
             "local extraction with local dependent on previously-defined local \
              (1)"
-            `Quick test_local_dependent_on_previously_defined_local;
-          test_case
-            "local extraction with local dependent on previously-defined local \
-             (2)"
-            `Quick test_local_not_ground_because_of_ordering;
+            `Quick test_local_dependent_on_previously_defined_local
         ] );
       ( "ast formatting ",
         [ test_case "identifier compilation" `Quick test_identifier ] );