2 files changed, 38 insertions, 20 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 0b22baf..cfafd8d 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -43,50 +43,63 @@ let rec uniq = function
       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
+type analysis_result =
+  { bound : string list;
+    free  : string list;
+    used  : string list }
+
+let extract_variables (env : environ) (q : query) : analysis_result =
+  let empty = { bound = []; free = []; used = [] } in
+  let combine =
+    fun x y ->
+      match (x, y) with
+      | ({ bound = bound';  free = free';  used = used' },
+         { bound = bound''; free = free''; used = used'' }) ->
+        { bound = bound' @ bound'';
+          free = free' @ free'';
+          used = used' @ used'' } in
+  let rec handle_term env : term -> analysis_result = function
     (* If we see a variable in an assignment, it's bound. *)
     | Call ((Var "assign", []), [ Var name; rhs ]) ->
-      let (x, y) = handle_term env rhs in
-      (x, [name] @ y)
+      let { bound; free; used } = handle_term env rhs in
+      { bound = [name] @ bound; free; used }
     (* Variables that haven't previously been introduced to the scope are free. *)
     | Var name ->
       if Map.find env name |> Option.is_some then
-        ([], [])
-      else ([name], [])
+        { free = []; bound = []; used = [name] }
+      else { free = [name]; bound = []; used = [] }
     (* Otherwise, we'll recurse on the term. *)
     | Ref (target, args) ->
       let terms = [term_of_ref_target target] @ args in
-      List.map terms ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
+      List.map terms ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
     | Array xs | Set xs ->
-      List.map xs ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
+      List.map xs ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
     | Object kvs -> List.map ~f:snd kvs
                     |> List.map ~f:(handle_term env)
-                    |> List.fold ~init:([], []) ~f:combine
+                    |> List.fold ~init:empty ~f:combine
     | Call (x, xs) ->
-      List.map ((Ref x : term)::xs) ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
+      List.map ((Ref x : term)::xs) ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
     | ArrayComprehension _ -> failwith "unimplemented"
     | ObjectComprehension _ -> failwith "unimplemented"
     | SetComprehension _ -> failwith "unimplemented"
-    | _ -> ([], [])
+    | _ -> empty
   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
+      List.map q ~f:(handle_literal env') |> List.fold ~init:empty ~f:combine
+  and handle_literal env : literal -> analysis_result = function
     | Exists (exprs, _) ->
-      List.map exprs ~f:(handle_expr env) |> List.fold ~init:([], []) ~f:combine
+      List.map exprs ~f:(handle_expr env) |> List.fold ~init:empty ~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; bound; used } = List.map q ~f:(handle_literal env) |> List.fold ~init:empty ~f:combine in
   let free = List.filter free  ~f:(fun x -> List.mem bound ~equal:String.equal x |> not) in
-  (free |> uniq, bound |> uniq)
+  { free = free |> uniq; bound = bound |> uniq; used = used |> uniq }
 
-let extract_variables_wrapper (q : query) : (string list * string list) =
+let extract_variables_wrapper (q : query) : analysis_result =
   extract_variables toplevel q
 
 (* --- *)
@@ -189,7 +202,7 @@ and compile_literal env : literal -> (environ * lean_expression) = function
   | Negate (_expr, _) -> failwith "unimplemented"
 
 and compile_query env literals : (environ * lean_expression) =
-  let (free, bound) = extract_variables env literals in
+  let { free; bound; _ } = extract_variables env literals in
   let env' = List.map free ~f:(fun name -> (name, { lean_id = name; complete = true }))
              |> Environ.of_alist_exn
              |> Map.merge env ~f:(fun ~key:_key -> function
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.mli b/rego-proofs/compiler-rewrite/lib/compiler.mli
index bdbc730..58a9c14 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.mli
+++ b/rego-proofs/compiler-rewrite/lib/compiler.mli
@@ -3,5 +3,10 @@ open Lean
 
 type environ
 
-val extract_variables_wrapper : query -> (string list * string list)
+type analysis_result =
+  { bound : string list;
+    free  : string list;
+    used  : string list }
+
+val extract_variables_wrapper : query -> analysis_result
 val compile_package : package -> lean_expression list