1 files changed, 70 insertions, 27 deletions
diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index c68dd8a..51f0ef7 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -315,38 +315,81 @@ let compile_definition env (rules : rule list) =
     | None -> Bool true
   in
 
-  let subject = gensym "x" in
+  let subjects = match List.nth_exn rules 0 with
+    | { head = ((_, []), Assignment (_, _)); _ } |
+      { head = ((_, _), Set _); _ } ->
+      [gensym "x"]
+    | { head = ((_, [_]), Assignment (_, _)); _ } ->
+      [gensym "k"; gensym "v"]
+    | { head = ((_, _), Function (_, _, _)); _ } ->
+      [gensym "xs"; gensym "y"]
+    | { head = ((_, _), Assignment (_, _)); _ } ->
+      failwith "unreachable"
+  in
+
+  (* let compile_eq = function *)
+
   let conditional_arms =
-    List.mapi conditional_rules ~f:(fun i ->
-        function
-        | { head = (_, Assignment (_, term)); body; _ } ->
-          let eq : term = Call ((Var "eq", []), [Var subject; term]) in
-          let env' = [(subject, { lean_id = subject; kind = Complete; assigned = true })]
-                     |> extend_environment env in
-          let (_, subject_constraint) = compile_query env' [Expr (Term eq, None)] in
-          let (_, body_form) = compile_query env body in
-          (body_form,
-           (Printf.sprintf "rule%d" i,
-            [subject_constraint; body_form],
-            [Identifier subject]))
-        | _ -> failwith "unimplemented") in
+    List.mapi conditional_rules ~f:(fun i rule ->
+        let env' =
+          List.map subjects
+            ~f:(fun subject ->
+                (subject, { lean_id = subject; kind = Complete; assigned = true }))
+          |> extend_environment env
+        in
+
+        let eq : term list = match rule with
+          | { head = ((_, []), Assignment (_, term)); _ } | { head = (_, Set (term)); _ } ->
+            [Call ((Var "eq", []), [Var (List.nth_exn subjects 0); term])]
+          | { head = ((_, [key]), Assignment (_, value)); _ } ->
+            [Call ((Var "eq", []), [Var (List.nth_exn subjects 0); key]);
+             Call ((Var "eq", []), [Var (List.nth_exn subjects 1); value])]
+          (* | { head = ((_, _), Function (_, args, term)); _ } -> *)
+          (*   [Call ((Var "eq", []), [Var (List.nth_exn subjects 0); Array args]); *)
+          (*    Call ((Var "eq", []), [Var (List.nth_exn subjects 1); term])] *)
+          | _ -> failwith "unimplemented"
+        in
+
+        let query = List.map eq ~f:(fun clause -> Expr (Term clause, None)) in
+        let (_, subject_constraint) = compile_query env' query in
+        let { body; _ } = rule in
+        let (_, body_form) = compile_query env body in
+        (body_form,
+         (Printf.sprintf "rule%d" i,
+          [subject_constraint; body_form],
+          List.map subjects ~f:(fun x -> Identifier x)))) in
   let default_arms =
-    List.map default_rules ~f:(function
-        | { head = (_, Assignment (_, term)); _ } ->
-          let eq : term = Call ((Var "eq", []), [Var subject; term]) in
-          let env' = [(subject, { lean_id = subject; kind = Complete; assigned = true })]
-                     |> extend_environment env in
-          let (_, subject_constraint) = compile_query env' [Expr (Term eq, None)] in
-          let other_rule_bodies = List.map ~f:fst conditional_arms in
-          let body_form = Application (Identifier "Not", [fold_or other_rule_bodies]) in
-          ("rule_default",
-           [subject_constraint; body_form],
-           [Identifier subject])
-        | _ -> failwith "unimplemented") in
+    List.map default_rules ~f:(fun rule ->
+        let env' =
+          List.map subjects
+            ~f:(fun subject ->
+                (subject, { lean_id = subject; kind = Complete; assigned = true }))
+          |> extend_environment env
+        in
+
+        let eq : term list = match rule with
+          | { head = ((_, []), Assignment (_, term)); _ } | { head = (_, Set (term)); _ } ->
+            [Call ((Var "eq", []), [Var (List.nth_exn subjects 0); term])]
+          | { head = ((_, [key]), Assignment (_, value)); _ } ->
+            [Call ((Var "eq", []), [Var (List.nth_exn subjects 0); key]);
+             Call ((Var "eq", []), [Var (List.nth_exn subjects 1); value])]
+          (* | { head = ((_, _), Function (_, args, term)); _ } -> *)
+          (*   [Call ((Var "eq", []), [Var (List.nth_exn subjects 0); Array args]); *)
+          (*    Call ((Var "eq", []), [Var (List.nth_exn subjects 1); term])] *)
+          | _ -> failwith "unimplemented"
+        in
+
+        let query = List.map eq ~f:(fun clause -> Expr (Term clause, None)) in
+        let (_, subject_constraint) = compile_query env' query in
+        let other_rule_bodies = List.map ~f:fst conditional_arms in
+        let body_form = Application (Identifier "Not", [fold_or other_rule_bodies]) in
+        ("rule_default",
+         [subject_constraint; body_form],
+         List.map subjects ~f:(fun x -> Identifier x))) in
   let form =
     Inductive
       (string_of_reference variable_being_defined,
-       [([subject], "Rego.Any")],
+       [(subjects, "Rego.Any")],
        List.append (List.map ~f:snd conditional_arms) default_arms)
   in (env, form)