Initial ast walking example

2 files changed, 34 insertions, 5 deletions

diff --git a/rego-proofs/compiler-rewrite/bin/dune b/rego-proofs/compiler-rewrite/bin/dune
index 5dd24b4..941cf84 100644
--- a/rego-proofs/compiler-rewrite/bin/dune
+++ b/rego-proofs/compiler-rewrite/bin/dune
@@ -3,7 +3,7 @@
  (name main)
  (libraries core compiler_rewrite unix yojson)
  (preprocess
-   (pps ppx_deriving.show)))
+   (pps ppx_deriving.show ppx_fields_conv)))
 
 (env
   (dev
diff --git a/rego-proofs/compiler-rewrite/bin/main.ml b/rego-proofs/compiler-rewrite/bin/main.ml
index 8d91a2e..c66fd9b 100644
--- a/rego-proofs/compiler-rewrite/bin/main.ml
+++ b/rego-proofs/compiler-rewrite/bin/main.ml
@@ -58,10 +58,10 @@ type rule_head = rule_head_lhs * rule_head_rhs
 type rule =
   { default : bool;
     head : rule_head;
-    body: literal list;
+    body: query;
     otherwise : rule option;
   }
-[@@deriving show]
+[@@deriving fields, show]
 
 type import =
   { alias : string option;
@@ -74,7 +74,33 @@ type package =
     imports: import list;
     policy : rule list;
   }
-[@@deriving show]
+[@@deriving fields, show]
+
+module Environ = Map.Make(String);;
+
+let extract_local_variables (q : query) (scope : string Environ.t) : (string list) =
+  let handle_term = function
+    | Call (((Var "assign"), []), [(Var name); _]) ->
+       if (Map.find scope name) |> Option.is_none then [name] else []
+    | _ -> []
+    (* | Var name -> if (Map.find scope name) |> Option.is_none then [name] else [] *)
+    (* | Array of term list *)
+    (* | Object of (object_key * term) list *)
+    (* | Set of term list *)
+  in
+  let handle_expr = function
+    | Term t -> handle_term t
+    | Every ((Some _key), _value, _, _q) -> failwith "unimplemented"
+    | Every (None, _value, _, _q) -> failwith "unimplemented"
+  in
+  let handle_literal = function
+    | Exists (exprs, _) -> List.map ~f:handle_expr exprs |> List.concat
+    | Expr (expr, _) -> handle_expr expr
+    | Negate (expr, _) -> handle_expr expr
+  in
+  List.map q ~f:handle_literal |> List.concat
+
+(* --- *)
 
 let is_null : Yojson.Basic.t -> bool = function
   | `Null -> true
@@ -251,4 +277,7 @@ let () =
   let file_name = Array.get (Sys.get_argv ()) 1 in
   let policy_code = In_channel.read_all file_name in
   let (stdout, _) = invoke ("/home/jakob/.local/bin/opa parse -f json /dev/stdin") ~stdin:policy_code ~env:[| |] in
-  Yojson.Basic.from_string stdout |> parse_tree |> show_package |> print_endline
+  let ast = Yojson.Basic.from_string stdout |> parse_tree in
+  match ast with
+  | { policy = [{ body = body; _ }]; _ } -> (extract_local_variables body Environ.empty) |> List.iter ~f:print_endline
+  | _ -> print_endline "failed"