Minor refactoring

1 files changed, 80 insertions, 86 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 51f0ef7..bd9ef14 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -7,6 +7,19 @@ type definition_kind = Complete
                      | SetRule
                      | ObjectRule
                      | FunctionRule
+[@@deriving eq, show]
+
+let definition_kind_of_rule = function
+  | { head = ((_, []), Assignment (_, _)); _ } ->
+    ScalarRule
+  | { head = ((_, _), Set _); _ } ->
+    SetRule
+  | { head = ((_, [_]), Assignment (_, _)); _ } ->
+    ObjectRule
+  | { head = ((_, _), Function (_, _, _)); _ } ->
+    FunctionRule
+  | { head = ((_, _), Assignment (_, _)); _ } ->
+    failwith "unreachable"
 
 type variable = { lean_id : string;
                   kind : definition_kind;
@@ -23,6 +36,11 @@ let extend_environment (env : environ) (vars : (string * variable) list) =
       | `Left v -> Some v
       | `Both (_, v) -> Some v)
 
+let toplevel = Environ.of_alist_exn [
+    ("eq",    { lean_id = "eq";    kind = Complete; assigned = true });
+    ("equal", { lean_id = "equal"; kind = Complete; assigned = true })
+  ]
+
 let term_of_object_key : object_key -> term = function
   | Scalar s -> Scalar s
   | Ref r    -> Ref r
@@ -45,17 +63,12 @@ let map_with_environment_updates env xs fn =
       let env', y = fn env x in
       (env', ys @ [ y ]))
 
-let toplevel = Environ.of_alist_exn [
-    ("eq",    { lean_id = "eq";    kind = Complete; assigned = true });
-    ("equal", { lean_id = "equal"; kind = Complete; assigned = true })
-  ]
-
-let rec uniq = function
+let rec uniq eq = function
   | [] -> []
   | hd::tl ->
-    if List.mem tl hd ~equal:String.equal then
-      uniq tl
-    else hd :: uniq tl
+    if List.mem tl hd ~equal:eq then
+      uniq eq tl
+    else hd :: uniq eq tl
 
 type analysis_result =
   { bound : string list;
@@ -116,8 +129,12 @@ let extract_variables (env : environ) (q : query) : analysis_result =
   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 = free |> uniq; bound = bound |> uniq; used = used |> uniq }
+  let uniq = uniq String.equal in
+  { free  = free  |> uniq;
+    bound = bound |> uniq;
+    used  = used  |> uniq }
 
+(* This is solely for the test suite. *)
 let extract_variables_wrapper (q : query) : analysis_result =
   extract_variables toplevel q
 
@@ -130,6 +147,8 @@ let gensym : string -> string =
     counter := !counter + 1 ;
     Printf.sprintf "%s__%d" base number
 
+(* This breaks a nested rule, such as one defining `w.x.y.z`, into as many as
+   needed such that every rule is of the form `x["y"] := _`. *)
 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 ->
@@ -289,68 +308,62 @@ and compile_query env literals : (environ * lean_expression) =
      Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
              Application (Identifier "And", compiled_terms)))
 
-(* TODO: Probably get rid of this once we deal with nested definitions. *)
-let string_of_reference : reference -> string = function
-  | (Var name, _) -> name
-  | _ -> failwith "unimplemented"
-
 let compile_definition env (rules : rule list) =
   assert (not (List.is_empty rules));
 
-  (* TODO: Assert that all rules are of the same kind. *)
+  (* Assert that all rules are of the same kind. *)
+  assert (1 = (List.map rules ~f:definition_kind_of_rule
+               |> uniq equal_definition_kind
+               |> List.length));
 
+  (* Rules are grouped based on their head, so we can determine which variable
+     we're defining by looking at any arbitrarily-chosen rule. It's easy enough
+     to get the first. *)
   let variable_being_defined =
-    match List.nth_exn rules 0 with
+    match List.hd_exn rules with
     | { head = (rf, _); _ } -> rf
   in
 
   let (default_rules, conditional_rules) =
-    List.partition_tf rules ~f:(function | { default = true; _ } -> true | _ -> false)
+    List.partition_tf rules ~f:(function
+        | { default = true; _ } -> true
+        | _ -> false)
   in
-  assert (List.length default_rules <= 1); (* There can be only one default rule. *)
+  (* There can be only one default rule. *)
+  assert (List.length default_rules <= 1);
 
-  let fold_or props =
-    match List.reduce props ~f:(fun x y -> Application (Identifier "Or", [y; x])) with
-    | Some x -> x
-    | None -> Bool true
-  in
-
-  let subjects = match List.nth_exn rules 0 with
-    | { head = ((_, []), Assignment (_, _)); _ } |
-      { head = ((_, _), Set _); _ } ->
+  let subjects = match List.nth_exn rules 0 |> definition_kind_of_rule with
+    | ScalarRule | SetRule ->
       [gensym "x"]
-    | { head = ((_, [_]), Assignment (_, _)); _ } ->
+    | ObjectRule ->
       [gensym "k"; gensym "v"]
-    | { head = ((_, _), Function (_, _, _)); _ } ->
+    | FunctionRule ->
       [gensym "xs"; gensym "y"]
-    | { head = ((_, _), Assignment (_, _)); _ } ->
+    | Complete ->
       failwith "unreachable"
   in
+  let env' =
+    List.map subjects
+      ~f:(fun subject ->
+          (subject, { lean_id = subject; kind = Complete; assigned = true }))
+    |> extend_environment env
+  in
 
-  (* let compile_eq = function *)
+  let mk_eq : rule -> term list = function
+    | { 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 conditional_arms =
     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 query = mk_eq rule |> List.map ~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
@@ -360,33 +373,27 @@ let compile_definition env (rules : rule list) =
           List.map subjects ~f:(fun x -> Identifier x)))) in
   let default_arms =
     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 query = mk_eq rule |> List.map ~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
+        let body_form =
+          let fold_or props =
+            match List.reduce props ~f:(fun x y -> Application (Identifier "Or", [y; x])) with
+            | Some x -> x
+            | None -> Bool true
+          in
+          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 =
+    let string_of_reference : reference -> string = function
+      | (Var name, []) -> name
+      | (Var name, [_]) -> name
+      | _ -> failwith "unreachable"
+    in
     Inductive
       (string_of_reference variable_being_defined,
        [(subjects, "Rego.Any")],
@@ -425,16 +432,3 @@ let compile_package = function
     [Comment (show_reference path |> String.escaped);
      Comment (List.map ~f:show_import imports |> String.concat |> String.escaped)]
     @ forms
-
-(* let subscript_of_number n = *)
-(*   string_of_int n *)
-(*   |> Str.global_replace (Str.regexp "0") "₀" *)
-(*   |> Str.global_replace (Str.regexp "1") "₁" *)
-(*   |> Str.global_replace (Str.regexp "2") "₂" *)
-(*   |> Str.global_replace (Str.regexp "3") "₃" *)
-(*   |> Str.global_replace (Str.regexp "4") "₄" *)
-(*   |> Str.global_replace (Str.regexp "5") "₅" *)
-(*   |> Str.global_replace (Str.regexp "6") "₆" *)
-(*   |> Str.global_replace (Str.regexp "7") "₇" *)
-(*   |> Str.global_replace (Str.regexp "8") "₈" *)
-(*   |> Str.global_replace (Str.regexp "9") "₉" *)