Add support for nested definitions

3 files changed, 52 insertions, 10 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index ecc7162..5c89807 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -10,6 +10,13 @@ type variable = { lean_id : string;
 module Environ = Map.Make (String)
 type environ = variable Environ.t
 
+let extend_environment (env : environ) (vars : (string * variable) list) =
+  Environ.of_alist_exn vars
+  |> Map.merge env ~f:(fun ~key:_key -> function
+      | `Right v -> Some v
+      | `Left v -> Some v
+      | `Both (_, v) -> Some v)
+
 let term_of_object_key : object_key -> term = function
   | Scalar s -> Scalar s
   | Ref r    -> Ref r
@@ -117,6 +124,18 @@ let gensym : string -> string =
     counter := !counter + 1 ;
     Printf.sprintf "%s__%d" base number
 
+let rec transform_rule_with_ref_head : rule -> rule list = function
+  (* Base case - this is a simple reference that we can deal with.  *)
+  | { head = ((_, []), _); _ } | { head = ((_, [_]), _); _ } as rule ->
+    [rule]
+  | { head = ((tgt, xs), rhs); _ } as rule  ->
+    let tail = List.last_exn xs in
+    let name = gensym "anonymous" in
+    [{ rule with head = ((Var name, [tail]), rhs) }] @
+    transform_rule_with_ref_head { rule with head = ((tgt, List.drop_last_exn xs), Assignment (Assign, Var name)) }
+
+(* --- *)
+
 let rec compile_term env : term -> lean_expression =
   function
   | Scalar s ->
@@ -216,14 +235,11 @@ and compile_literal env : literal -> (environ * lean_expression) = function
 
 and compile_query env literals : (environ * lean_expression) =
   let { free; bound; _ } = extract_variables env literals in
-  let env' = List.map free ~f:(fun name -> (name, { lean_id = gensym name; complete = true; assigned = true }))
-             |> List.append
-               (List.map bound ~f:(fun name -> (name, { lean_id = gensym name; complete = true; assigned = false })))
-             |> Environ.of_alist_exn
-             |> Map.merge env ~f:(fun ~key:_key -> function
-                 | `Right v -> Some v
-                 | `Left v -> Some v
-                 | `Both (_, v) -> Some v) in
+  let env' =
+    List.append
+      (List.map free  ~f:(fun name -> (name, { lean_id = gensym name; complete = true; assigned = true })))
+      (List.map bound ~f:(fun name -> (name, { lean_id = gensym name; complete = true; assigned = false })))
+    |> extend_environment env in
   let (env'', subterms) = map_with_environment_updates env' literals compile_literal in
   let compiled_terms = List.map subterms
       ~f:(fun x -> Application (Identifier "Ne", [x; Application (Identifier "Rego.Any.bool", [Bool false])])) in
@@ -258,9 +274,9 @@ let compile_definition env (rules : rule list) =
        List.mapi rules ~f:(fun i ->
            function
            | { head = (_, Assignment (_, term)); body; _ } ->
-             let form = compile_term toplevel term in
+             let form = compile_term env term in
              (Printf.sprintf "rule%d" i,
-              compile_query toplevel body |> snd,
+              compile_query env body |> snd,
               [form])
            | _ -> failwith "unimplemented"))
   in (env, form)
@@ -288,7 +304,16 @@ let compile_policy env (rules : rule list) =
 
 let compile_package = function
   | { path; imports; policy } ->
+    let policy = List.map ~f:transform_rule_with_ref_head policy |> List.concat in
+
     let env = toplevel in
+    let env =
+      List.map policy ~f:(function
+          | { head = ((Var name, _), _); _ } ->
+            (name, { lean_id = name ; complete = false; assigned = true })
+          | _ -> failwith "policy rule head is not a variable")
+      |> extend_environment env 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)]
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.mli b/rego-proofs/compiler-rewrite/lib/compiler.mli
index 58a9c14..0bef961 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.mli
+++ b/rego-proofs/compiler-rewrite/lib/compiler.mli
@@ -9,4 +9,5 @@ type analysis_result =
     used  : string list }
 
 val extract_variables_wrapper : query -> analysis_result
+val transform_rule_with_ref_head : rule -> rule 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 cbb7b04..847f8c9 100644
--- a/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
+++ b/rego-proofs/compiler-rewrite/test/compiler_rewrite.ml
@@ -3,6 +3,7 @@ open Compiler_rewrite
 
 module To_test = struct
   let extract_variables = Compiler.extract_variables_wrapper
+  let transform_ref_head = Compiler.transform_rule_with_ref_head
   let fmt = Lean.fmt_lean_expression
 end
 
@@ -95,6 +96,16 @@ let test_identifier () =
   let ast = Identifier "test" in
   Alcotest.(check string) "expected string" "test" (To_test.fmt ast)
 
+let test_transform_simple_1 () =
+  let open Compiler_rewrite.Ast in
+  let rule = { head = ((Var "fruit", [Var "apple"; Var "seeds"; Var "density"]), Assignment (Assign, (Scalar (Number 12.))));
+               default = false;
+               body = [Expr (Term (Scalar (Bool true)), None)];
+               otherwise = None } in
+  Alcotest.(check int) "expected rules"
+    2
+    (To_test.transform_ref_head rule |> List.length)
+
 let () =
   let open Alcotest in
   run "Compiler"
@@ -117,6 +128,11 @@ let () =
              (1)"
             `Quick test_local_dependent_on_previously_defined_local
         ] );
+      ( "nested definitions",
+        [
+          test_case "local nested definitions (1)" `Quick
+            test_transform_simple_1;
+        ] ) ;
       ( "ast formatting ",
         [ test_case "identifier compilation" `Quick test_identifier ] );
     ]