Re-do first pass variable identification

4 files changed, 132 insertions, 83 deletions

diff --git a/rego-proofs/compiler-rewrite/bin/main.ml b/rego-proofs/compiler-rewrite/bin/main.ml
index 933d97a..b5f04fa 100644
--- a/rego-proofs/compiler-rewrite/bin/main.ml
+++ b/rego-proofs/compiler-rewrite/bin/main.ml
@@ -8,6 +8,5 @@ let () =
   let policy_code = In_channel.read_all file_name in
   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 a0f521b..78882c1 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -9,51 +9,6 @@ type variable = { lean_id : string;
 module Environ = Map.Make (String)
 type environ = variable Environ.t
 
-let rec term_ground (env : environ) (term : term) : bool =
-  match term with
-  | Var name -> Map.find env name |> Option.is_some
-  | Ref (Var name, _) -> Map.find env name |> Option.is_some
-  | Call (_, subterms) ->
-      List.reduce_exn ~f:( && ) (List.map ~f:(term_ground env) subterms)
-  | _ -> true
-
-let map_with_environment_updates env xs fn =
-  List.fold xs ~init:(env, []) ~f:(fun (env, ys) x ->
-      let env', y = fn env x in
-      (env', List.append ys [ y ]))
-
-let extract_local_variables (env : environ) (q : query) : string list =
-  let handle_term env : term -> (environ * string list) = function
-    | Call ((Var func, []), [ Var name; rhs ])
-      when String.equal func "assign" || String.equal func "eq" ->
-        if Map.find env name |> Option.is_none && term_ground env rhs then
-          match Map.add env ~key:name ~data:{ lean_id = name; complete = true} with
-          | `Duplicate -> failwith "should shadow?"
-          | `Ok env' -> (env', [ name ])
-        else (env, [])
-    | _ -> (env, [])
-  in
-  let handle_expr env = function
-    | Term t -> handle_term env t
-    | Every (Some _key, _value, _, _q) -> failwith "unimplemented"
-    | Every (None, _value, _, _q) -> failwith "unimplemented"
-  in
-  let handle_literal env : literal -> (environ * string list) = function
-    | Exists (exprs, _) ->
-        let env', ys = map_with_environment_updates env exprs handle_expr in
-        (env', ys |> List.concat)
-    | Expr (expr, _) -> handle_expr env expr
-    | Negate (expr, _) -> handle_expr env expr
-  in
-
-  let _env', identifiers = map_with_environment_updates env q handle_literal in
-  List.concat identifiers
-
-let extract_local_variables_wrapper (q : query) : string list =
-  extract_local_variables Environ.empty q
-
-(* --- *)
-
 let term_of_object_key : object_key -> term = function
   | Scalar s -> Scalar s
   | Ref r    -> Ref r
@@ -71,6 +26,62 @@ let term_of_ref_target : ref_target -> term = function
 
 let literal_of_expr exp : literal = Expr (exp, None)
 
+let map_with_environment_updates env xs fn =
+  List.fold xs ~init:(env, []) ~f:(fun (env, ys) x ->
+      let env', y = fn env x in
+      (env', ys @ [ y ]))
+
+let toplevel = Environ.of_alist_exn [
+    ("eq",    { lean_id = "eq";    complete = true });
+    ("equal", { lean_id = "equal"; complete = true })
+  ]
+
+let rec uniq = function
+  | [] -> []
+  | hd::tl ->
+    if List.mem tl hd ~equal:String.equal then
+      uniq tl
+    else hd :: uniq tl
+
+let extract_variables (env : environ) (q : query) : (string list * string list) =
+  let combine = fun (x, y) (x', y') -> (x @ x', y @ y') in
+  let rec handle_term env : term -> (string list * string list) = function
+    | Call ((Var "assign", []), [ Var name; rhs ]) -> let (x, y) = handle_term env rhs in (x, [name] @ y)
+    | Var name -> if Map.find env name |> Option.is_some then ([], []) else ([name], [])
+    | Ref (target, args) ->
+      let terms = [term_of_ref_target target] @ args in
+      List.map terms ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
+    | Array xs | Set xs -> List.map xs ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
+    | Object kvs -> List.map ~f:snd kvs
+                    |> List.map ~f:(handle_term env)
+                    |> List.fold ~init:([], []) ~f:combine
+    | Call (x, xs) ->
+      List.map ((Ref x : term)::xs) ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
+    | ArrayComprehension _ -> failwith "unimplemented"
+    | ObjectComprehension _ -> failwith "unimplemented"
+    | SetComprehension _ -> failwith "unimplemented"
+    | _ -> ([], [])
+  and handle_expr env = function
+    | Term t -> handle_term env t
+    | Every (Some _key, _value, _, _q) -> failwith "unimplemented"
+    | Every (None, value, _, q) ->
+      let env' = Map.set env ~key:value ~data:({ lean_id = "test"; complete = true }) in
+      List.map q ~f:(handle_literal env') |> List.fold ~init:([], []) ~f:combine
+  and handle_literal env : literal -> (string list * string list) = function
+    | Exists (exprs, _) ->
+      List.map exprs ~f:(handle_expr env) |> List.fold ~init:([], []) ~f:combine
+    | Expr (expr, _) -> handle_expr env expr
+    | Negate (expr, _) -> handle_expr env expr
+  in
+  let (free, bound) = List.map q ~f:(handle_literal env) |> List.fold ~init:([], []) ~f:combine in
+  let free = List.filter free  ~f:(fun x -> List.mem bound ~equal:String.equal x |> not) in
+  (free |> uniq, bound |> uniq)
+
+let extract_variables_wrapper (q : query) : (string list * string list) =
+  extract_variables toplevel q
+
+(* --- *)
+
 type compile_term_result =
   { form : lean_expression;
     toplevel : lean_expression list;
@@ -225,7 +236,7 @@ let string_of_reference : reference -> string = function
   | (Var name, _) -> name
   | _ -> failwith "unimplemented"
 
-let compile_definition (rules : rule list) =
+let compile_definition env (rules : rule list) =
   assert (not (List.is_empty rules));
 
   let variable_being_defined =
@@ -233,21 +244,23 @@ let compile_definition (rules : rule list) =
     | { 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 form =
+    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"))
+  in (env, form)
 
-let compile_policy (rules : rule list) =
+let compile_policy env (rules : rule list) =
   let group accessor l =
     let rec grouping acc = function
       | [] -> acc
@@ -266,8 +279,25 @@ let compile_policy (rules : rule list) =
   (*       ) *)
   (* in *)
 
-  List.map by_reference ~f:compile_definition
+  map_with_environment_updates env by_reference compile_definition
 
 let compile_package = function
-  | { policy; _ } ->
-    [(* Comment path; Comment imports *)compile_policy policy]
+  | { path; imports; policy } ->
+    let env = Environ.empty in
+    let (_env', forms) = compile_policy env policy in
+    [Comment (show_reference path |> String.escaped);
+     Comment (List.map ~f:show_import imports |> String.concat |> String.escaped)]
+    @ forms
+
+(* let subscript_of_number n = *)
+(*   string_of_int n *)
+(*   |> Str.global_replace (Str.regexp "0") "₀" *)
+(*   |> Str.global_replace (Str.regexp "1") "₁" *)
+(*   |> Str.global_replace (Str.regexp "2") "₂" *)
+(*   |> Str.global_replace (Str.regexp "3") "₃" *)
+(*   |> Str.global_replace (Str.regexp "4") "₄" *)
+(*   |> Str.global_replace (Str.regexp "5") "₅" *)
+(*   |> Str.global_replace (Str.regexp "6") "₆" *)
+(*   |> Str.global_replace (Str.regexp "7") "₇" *)
+(*   |> Str.global_replace (Str.regexp "8") "₈" *)
+(*   |> Str.global_replace (Str.regexp "9") "₉" *)
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.mli b/rego-proofs/compiler-rewrite/lib/compiler.mli
index 66bd4dc..bdbc730 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.mli
+++ b/rego-proofs/compiler-rewrite/lib/compiler.mli
@@ -3,5 +3,5 @@ open Lean
 
 type environ
 
-val extract_local_variables_wrapper : query -> string list
-val compile_package : package -> lean_expression list list
+val extract_variables_wrapper : query -> (string list * string list)
+val compile_package : package -> lean_expression list
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 ] );