6 files changed, 23 insertions, 31 deletions
diff --git a/rego-proofs/compiler-rewrite/bin/dune b/rego-proofs/compiler-rewrite/bin/dune
index 7990455..a53c710 100644
--- a/rego-proofs/compiler-rewrite/bin/dune
+++ b/rego-proofs/compiler-rewrite/bin/dune
@@ -1,5 +1,5 @@
 (executable
  (public_name compiler_rewrite)
  (name main)
- (libraries core compiler_rewrite yojson))
+ (libraries core compiler_rewrite))
 
diff --git a/rego-proofs/compiler-rewrite/bin/main.ml b/rego-proofs/compiler-rewrite/bin/main.ml
index 3390ac7..fb9f5e5 100644
--- a/rego-proofs/compiler-rewrite/bin/main.ml
+++ b/rego-proofs/compiler-rewrite/bin/main.ml
@@ -1,21 +1,10 @@
 open Core
-open Compiler_rewrite.Compiler
-
-let invoke ?(stdin = "") ?(env = [||]) cmd =
-  let ic, oc, ec = UnixLabels.open_process_full cmd ~env in
-  Out_channel.output_string oc stdin;
-  Out_channel.close oc;
-  (In_channel.input_all ic, In_channel.input_all ec)
+open Compiler_rewrite.Ast
 
 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
-  let ast = Yojson.Basic.from_string stdout |> parse_tree in
-  match ast with
+  match ast_of_string policy_code with
   | { policy = [ { body; _ } ]; _ } ->
       extract_local_variables_wrapper body |> List.iter ~f:print_endline
   | _ -> print_endline "failed"
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/ast.ml
index 4af14d1..884f2ff 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/ast.ml
@@ -295,3 +295,18 @@ let parse_tree tree =
       List.map ~f:parse_import (tree |> member "imports" |> to_list_or_empty);
     policy = List.map ~f:parse_rule (tree |> member "rules" |> to_list_or_empty);
   }
+
+(* --- *)
+
+let invoke ?(stdin = "") ?(env = [||]) cmd =
+  let ic, oc, ec = UnixLabels.open_process_full cmd ~env in
+  Out_channel.output_string oc stdin;
+  Out_channel.close oc;
+  (In_channel.input_all ic, In_channel.input_all ec)
+
+let ast_of_string policy_code =
+  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
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.mli b/rego-proofs/compiler-rewrite/lib/ast.mli
index b7b14dc..b373b0d 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.mli
+++ b/rego-proofs/compiler-rewrite/lib/ast.mli
@@ -58,4 +58,4 @@ type import = { alias : string option; path : reference }
 type package = { path : reference; imports : import list; policy : rule list }
 
 val extract_local_variables_wrapper : query -> string list
-val parse_tree : Yojson.Basic.t -> package
+val ast_of_string : string -> package
diff --git a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
index e9484a6..c21d157 100644
--- a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
+++ b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
@@ -1,7 +1,8 @@
 open Core
+open Compiler_rewrite
 
 module To_test = struct
-  let extract_locals = Compiler_rewrite.Compiler.extract_local_variables_wrapper
+  let extract_locals = Ast.extract_local_variables_wrapper
 end
 
 let strings =
@@ -9,21 +10,8 @@ let strings =
     (Format.pp_print_list Format.pp_print_string)
     (List.equal String.equal)
 
-let invoke ?(stdin = "") ?(env = [||]) cmd =
-  let ic, oc, ec = UnixLabels.open_process_full cmd ~env in
-  Out_channel.output_string oc stdin;
-  Out_channel.close oc;
-  (In_channel.input_all ic, In_channel.input_all ec)
-
-let rego_to_ast policy_code =
-  let stdout, _ =
-    invoke "/home/jakob/.local/bin/opa parse -f json /dev/stdin"
-      ~stdin:policy_code ~env:[||]
-  in
-  Yojson.Basic.from_string stdout |> Compiler_rewrite.Compiler.parse_tree
-
 let extract_query policy =
-  match rego_to_ast policy with
+  match Ast.ast_of_string policy with
   | { policy = [ { body; _ } ]; _ } -> body
   | _ -> failwith "failed to parse"
 
diff --git a/rego-proofs/compiler-rewrite/test/dune b/rego-proofs/compiler-rewrite/test/dune
index 273167d..63a9a47 100644
--- a/rego-proofs/compiler-rewrite/test/dune
+++ b/rego-proofs/compiler-rewrite/test/dune
@@ -1,3 +1,3 @@
 (test
  (name compiler_rewrite)
- (libraries alcotest compiler_rewrite core unix yojson))
+ (libraries alcotest compiler_rewrite core))