Basic expression simplification

3 files changed, 49 insertions, 5 deletions

diff --git a/rego-proofs/compiler-rewrite/lib/compiler.ml b/rego-proofs/compiler-rewrite/lib/compiler.ml
index 4756304..5f8f797 100644
--- a/rego-proofs/compiler-rewrite/lib/compiler.ml
+++ b/rego-proofs/compiler-rewrite/lib/compiler.ml
@@ -10,6 +10,7 @@ type definition_kind = Complete
                      | FunctionRule
 [@@deriving eq, show]
 
+
 let kind_and_args_of_rule = function
   | { head = ((_, []), Assignment (_, term)); _ } ->
     `ScalarRule term
@@ -182,6 +183,37 @@ let rec transform_rule_with_ref_head : rule -> rule list = function
 
 (* --- *)
 
+let rec simplify_expression = function
+  | Application (Identifier "And", xs) ->
+    let xs' = List.map ~f:simplify_expression xs
+              |> List.filter ~f:(function Bool true -> false | _ -> true)
+    in
+    (match xs' with
+     | [] -> Bool false
+     | [x] -> x
+     | xs' -> Application (Identifier "And", xs'))
+  | Application (Identifier "Eq", [Application (Identifier "Rego.Any.bool", [Bool a]);
+                                   Application (Identifier "Rego.Any.bool", [Bool b])]) ->
+    Bool (equal_bool a b)
+  | Application (Identifier "Ne", [Application (Identifier "Rego.Any.bool", [Bool a]);
+                                   Application (Identifier "Rego.Any.bool", [Bool b])]) ->
+    Bool (not (equal_bool a b))
+  | Application (Identifier "Eq", [Application (Identifier "Rego.Any.str", [String a]);
+                                   Application (Identifier "Rego.Any.str", [String b])]) ->
+    Bool (equal_string a b)
+  | Application (Identifier "Ne", [Application (Identifier "Rego.Any.str", [String a]);
+                                   Application (Identifier "Rego.Any.str", [String b])]) ->
+    Bool (not (equal_string a b))
+  | Application (Identifier "Eq", [Application (Identifier "Rego.Any.num", [Number a]);
+                                   Application (Identifier "Rego.Any.num", [Number b])]) ->
+    Bool (equal_float a b)
+  | Application (Identifier "Ne", [Application (Identifier "Rego.Any.num", [Number a]);
+                                   Application (Identifier "Rego.Any.num", [Number b])]) ->
+    Bool (not (equal_float a b))
+  | expr -> expr
+
+(* --- *)
+
 let rec compile_term env : term -> lean_expression =
   function
   | Scalar s ->
@@ -307,7 +339,8 @@ let rec compile_expr env : expr -> (environ * lean_expression) = function
             [(Application
                 ((Identifier "Membership.mem"),
                  [Tuple [Identifier key; Identifier value]; domain']));
-             query']))))
+             query']))
+       |> simplify_expression))
   | Every (None, value, domain, query) ->
     let domain' = compile_term env domain in
     (* Extend the environment with the value. *)
@@ -321,7 +354,8 @@ let rec compile_expr env : expr -> (environ * lean_expression) = function
             [(Application
                 ((Identifier "Membership.mem"),
                  [Identifier value; domain']));
-             query']))))
+             query']))
+       |> simplify_expression))
 
 and compile_literal env : literal -> (environ * lean_expression) = function
   | Exists (exprs, _) ->
@@ -338,7 +372,7 @@ and compile_query env literals : (environ * lean_expression) =
   let (env', forms) = map_with_environment_updates env literals compile_literal in
   let form = match forms with
     | [x] -> x
-    | xs -> Application (Identifier "And", xs) in
+    | xs -> Application (Identifier "And", xs) |> simplify_expression in
   (env', form)
 
 and compile_rule env subjects rule : (environ * lean_expression) =
@@ -447,9 +481,11 @@ and compile_rule env subjects rule : (environ * lean_expression) =
   (env'',
    if List.is_empty quantified_names then
      Application (Identifier "And", subject_constraints @ [form])
+     |> simplify_expression
    else
      Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
-             Application (Identifier "And", subject_constraints @ constraints @ [form])))
+             Application (Identifier "And", subject_constraints @ constraints @ [form])
+             |> simplify_expression))
 
 
 let compile_definition env (rules : rule list) =
diff --git a/rego-proofs/compiler-rewrite/lib/lean.ml b/rego-proofs/compiler-rewrite/lib/lean.ml
index 41abee3..b3526e3 100644
--- a/rego-proofs/compiler-rewrite/lib/lean.ml
+++ b/rego-proofs/compiler-rewrite/lib/lean.ml
@@ -13,8 +13,9 @@ type lean_expression =
   | Inductive of string * judgement list * (string * lean_expression list * lean_expression list) list
   | ForAll of judgement * lean_expression
   | Exists of judgement * lean_expression
+[@@deriving eq, show]
 
-and judgement = string list * string
+and judgement = string list * string [@@deriving eq, show]
 
 let fmt_judgement : judgement -> string = function
   | vars, t -> Printf.sprintf "(%s : %s)" (String.concat ~sep:" " vars) t
diff --git a/rego-proofs/compiler-rewrite/lib/lean.mli b/rego-proofs/compiler-rewrite/lib/lean.mli
index 990ab5f..fb58c72 100644
--- a/rego-proofs/compiler-rewrite/lib/lean.mli
+++ b/rego-proofs/compiler-rewrite/lib/lean.mli
@@ -14,4 +14,11 @@ type lean_expression =
 
 and judgement = string list * string
 
+val equal_lean_expression : lean_expression -> lean_expression -> bool
+val equal_judgement : judgement -> judgement -> bool
+val pp_lean_expression : Format.formatter -> lean_expression -> unit
+val show_lean_expression : lean_expression -> string
+val pp_judgement : Format.formatter -> judgement -> unit
+val show_judgement : judgement -> string
+
 val fmt_lean_expression : lean_expression -> string