Return side information when compiling terms

1 files changed, 70 insertions, 13 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index fbf5ceb..01e076d 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -2,8 +2,12 @@ open Ast
 open Core
 open Lean
 
+type variable = { lean_id : string;
+                  complete : bool;
+                }
+
 module Environ = Map.Make (String)
-type environ = string Environ.t
+type environ = variable Environ.t
 
 let rec term_ground (env : environ) (term : term) : bool =
   match term with
@@ -23,7 +27,7 @@ let extract_local_variables (env : environ) (q : query) : string list =
     | 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:name with
+          match Map.add env ~key:name ~data:{ lean_id = name; complete = true} with
           | `Duplicate -> failwith "should shadow?"
           | `Ok env' -> (env', [ name ])
         else (env, [])
@@ -65,30 +69,80 @@ let term_of_ref_target : ref_target -> term = function
   | SetComprehension (key, q)    -> SetComprehension (key, q)
   | Call (fn, args)              -> Call (fn, args)
 
-let rec compile_term env : term -> lean_expression = function
+type compile_term_result =
+  { form : lean_expression;
+    toplevel : lean_expression list;
+    new_locals : (string * string) list }
+
+let gensym : string -> string =
+  let counter = ref 0 in
+  fun (base : string) ->
+    let number = !counter in
+    counter := !counter + 1 ;
+    Printf.sprintf "%s__%d" base number
+
+let rec compile_term env : term -> compile_term_result =
+  let result_of expr =
+    { form = expr;
+      toplevel = [];
+      new_locals = []; } in
+  let form = function | { form ; _ } -> form in
+  let toplevel = function | { toplevel ; _ } -> toplevel in
+  let new_locals = function | { new_locals ; _ } -> new_locals in
+
+  function
   | Scalar s ->
      (match s with
       | String s -> Application (Identifier "Rego.Any.str",  [String s])
       | Number x -> Application (Identifier "Rego.Any.num",  [Number x])
       | Bool b   -> Application (Identifier "Rego.Any.bool", [Bool b])
       | Null     -> Application (Identifier "Rego.Any.null", []))
-  | Var id -> Identifier (Map.find_exn env id)
+     |> result_of
+  | Var id ->
+     (match Map.find_exn env id with
+      | { lean_id; complete = true } -> Identifier lean_id |> result_of
+      | { lean_id; complete = false } ->
+        let anonymous_id = gensym "anonymous" in
+        { form       = Identifier anonymous_id;
+          toplevel   = [];
+          new_locals = [(anonymous_id, lean_id)]; })
   | Ref (target, terms) ->
     List.fold
       terms
       ~init:(compile_term env (term_of_ref_target target))
-      ~f:(fun x y -> (Application ((Identifier "Rego.Any.deref"), [x; compile_term env y])))
+      ~f:(fun x y ->
+          match x with
+          | { form = x'; toplevel = toplevel'; new_locals = new_locals' } ->
+            match compile_term env y with
+            | { form = y'; toplevel = toplevel''; new_locals = new_locals'' } ->
+              let toplevel = toplevel' @ toplevel'' in
+              let new_locals = new_locals' @ new_locals'' in
+              { form = (Application ((Identifier "Rego.Any.deref"), [x'; y']));
+                toplevel; new_locals })
   | Array xs ->
-     Application (Identifier "Rego.Any.array", List.map ~f:(compile_term env) xs)
+    let subterms = List.map ~f:(compile_term env) xs in
+    { form = Application (Identifier "Rego.Any.array", List.map ~f:form subterms);
+      toplevel = List.map ~f:toplevel subterms |> List.concat;
+      new_locals = List.map ~f:new_locals subterms |> List.concat }
   | Set xs ->
-     Application (Identifier "Rego.Any.set", List.map ~f:(compile_term env) xs)
+    let subterms = List.map ~f:(compile_term env) xs in
+    { form = Application  (Identifier "Rego.Any.set", List.map ~f:form subterms);
+      toplevel = List.map ~f:toplevel subterms |> List.concat;
+      new_locals = List.map ~f:new_locals subterms |> List.concat }
   | Object kvs ->
      let compile_kv = function
        | (key, value) ->
-          Tuple [compile_term env (term_of_object_key key);
-                 compile_term env value]
-     in
-     Application (Identifier "Rego.Any.object", List.map ~f:compile_kv kvs)
+         match compile_term env (term_of_object_key key) with
+         | { form = x'; toplevel = toplevel'; new_locals = new_locals' } ->
+           match compile_term env value with
+           | { form = y'; toplevel = toplevel''; new_locals = new_locals'' } ->
+             let toplevel = toplevel' @ toplevel'' in
+             let new_locals = new_locals' @ new_locals'' in
+             { form = Tuple [x'; y']; toplevel; new_locals } in
+     let subterms = List.map ~f:compile_kv kvs in
+     { form = Application  (Identifier "Rego.Any.object", List.map ~f:form subterms);
+       toplevel = List.map ~f:toplevel subterms |> List.concat;
+       new_locals = List.map ~f:new_locals subterms |> List.concat }
   | ArrayComprehension (_term, _query) ->
      failwith "unimplemented"
   | SetComprehension (_term, _query) ->
@@ -96,5 +150,8 @@ let rec compile_term env : term -> lean_expression = function
   | ObjectComprehension ((_key, _value), _query) ->
      failwith "unimplemented"
   | Call (fn, args) ->
-     let head = compile_term env (Ref fn) in
-     Application (head, List.map ~f:(compile_term env) args)
+    let head = compile_term env (Ref fn) in
+    let subterms = List.map ~f:(compile_term env) args in
+    { form = Application (form head, List.map ~f:form subterms);
+      toplevel = toplevel head @ (List.map ~f:toplevel subterms |> List.concat);
+      new_locals = new_locals head @ (List.map ~f:new_locals subterms |> List.concat) }