Better support for parsing bodies

1 files changed, 49 insertions, 35 deletions

diff --git a/rego-proofs/compiler-rewrite/bin/main.ml b/rego-proofs/compiler-rewrite/bin/main.ml
index 1bb28c1..609d608 100644
--- a/rego-proofs/compiler-rewrite/bin/main.ml
+++ b/rego-proofs/compiler-rewrite/bin/main.ml
@@ -13,6 +13,8 @@ type scalar = String of string
             | Bool of bool
             | Null
 [@@deriving show]
+and reference = ref_target * (ref_arg list)
+[@@deriving show]
 and ref_target = Var of string
                 | Array of term list
                 | Object of (object_item * object_item)
@@ -20,7 +22,7 @@ and ref_target = Var of string
                 | ArrayComprehension of term * query
                 | ObjectComprehension of object_item * query
                 | SetComprehension of term * query
-                | Call of (string list) * (expr list)
+                | Call of (string list) * (term list)
 [@@deriving show]
 and ref_arg = Scalar of scalar
              | Var of string
@@ -29,7 +31,7 @@ and ref_arg = Scalar of scalar
              | Set of term list
              | Ignore
 [@@deriving show]
-and object_item = Scalar of scalar | Ref of ref_target * (ref_arg list) | Var
+and object_item = Scalar of scalar | Ref of reference | Var
 [@@deriving show]
 and ref_item = Var of string
               | Array of term list
@@ -38,9 +40,9 @@ and ref_item = Var of string
               | ArrayComprehension of term * query
               | ObjectComprehension of object_item * query
               | SetComprehension of term * query
-              | Call of (string list) * (expr list)
+              | Call of (string list) * (term list)
 [@@deriving show]
-and term = Ref of ref_target * (ref_arg list)
+and term = Ref of reference
           | Var of string
           | Scalar of scalar
           | Array of term list
@@ -67,20 +69,16 @@ and infix_operator = BoolEq
                    | Unify
 [@@deriving show]
 and every_item = Term of term
-               | Call of (string list) * (expr list)
+               | Call of (string list) * (term list)
                | Infix of expr * infix_operator * expr
 [@@deriving show]
 and expr = Term of term
-         | Call of (string list) * (expr list)
+         | Call of reference * (term list)
          | Infix of expr * infix_operator * expr
          | Every of (string list) * every_item * query
          | Unary of expr (* Has to be unary negation *)
 [@@deriving show]
-and literal_head = Some of (term list) * expr
-            | Expr of expr
-            | Negate of expr
-[@@deriving show]
-and literal = literal_head * ((term * term) option) (* "with" modifier *)
+and literal = expr * ((term * term) option) (* "with" modifier *)
 [@@deriving show]
 and query = literal list
 [@@deriving show]
@@ -88,7 +86,7 @@ and query = literal list
 type assign_operator = Assign | Unify
 [@@deriving show]
 
-type rule_head_lhs = ref_target * (ref_arg list)
+type rule_head_lhs = reference
 [@@deriving show]
 type rule_head_rhs = Assignment of assign_operator * term
                    | Object of term * assign_operator * term
@@ -97,26 +95,22 @@ type rule_head_rhs = Assignment of assign_operator * term
 [@@deriving show]
 type rule_head = rule_head_lhs * rule_head_rhs
 [@@deriving show]
-
-type rule_body = Query of query | Literal of literal
-[@@deriving show]
-
 type rule =
   { default : bool;
     head : rule_head;
-    body: rule_body list;
+    body: literal list;
     otherwise : rule option;
   }
 [@@deriving show]
 
 type import =
   { alias : string option;
-    path : ref_target * (ref_arg list);
+    path : reference;
   }
 [@@deriving show]
 
 type package =
-  { path : ref_target * (ref_arg list);
+  { path : reference;
     imports: import list;
     policy : rule list;
   }
@@ -131,6 +125,8 @@ let to_list_or_empty : Yojson.Basic.t -> Yojson.Basic.t list = function
   | `List xs -> xs
   | _ -> failwith "Not a list or null"
 
+(* --- *)
+
 let parse_ref_target tree : ref_target =
   match (tree |> member "type" |> to_string) with
   | "var" -> Var (tree |> member "value" |> to_string)
@@ -146,15 +142,7 @@ let parse_ref tree =
   | target :: args -> (parse_ref_target target, List.map ~f:parse_ref_arg args)
   | _ -> failwith "empty reference"
 
-let parse_import_path tree =
-  match (tree |> member "type" |> to_string) with
-  | "ref" -> tree |> member "value" |> to_list |> parse_ref
-  | _ -> failwith "unimplemented"
-
-let parse_import tree =
-  { alias = tree |> member "alias" |> to_string_option;
-    path = tree |> member "path" |> parse_import_path
-  }
+(* --- *)
 
 let parse_term (tree : Yojson.Basic.t) : term =
   match tree |> member "type" |> to_string with
@@ -162,29 +150,55 @@ let parse_term (tree : Yojson.Basic.t) : term =
   | "number" -> Scalar (Number (tree |> member "value" |> to_number))
   | _ -> failwith "unimplemented"
 
+let parse_expr tree : expr =
+  print_endline (Yojson.Basic.to_string tree);
+  match tree |> member "terms" with
+  | `List (func :: args) ->
+     Call (func |> member "value" |> to_list |> parse_ref,
+           List.map ~f:parse_term args)
+  | `List _ -> failwith "todo"
+  | x -> Term (parse_term x)
+
 let parse_head tree =
   (* Variables are desugared to empty references. *)
   let lhs = tree |> member "ref" |> to_list |> parse_ref in
   let rhs =
     if tree |> member "assign" |> to_bool_or_false then
       Assignment (Assign, (tree |> member "value" |> parse_term))
+    (* TODO: Rest cases? *)
     else
-      failwith "unimplemented"
+      Assignment (Unify, (tree |> member "value" |> parse_term))
   in (lhs, rhs)
 
-let parse_body tree =
-  Literal (Expr (Term (tree |> member "terms" |> parse_term)), None)
-
-let parse_rule tree =
+let rec parse_rule tree =
   { default = tree |> member "default" |> to_bool_or_false;
     head = tree |> member "head" |> parse_head;
-    body = List.map ~f:parse_body (tree |> member "body" |> to_list_or_empty);
-    otherwise = None;
+    body = List.map ~f:(fun x -> ((parse_expr x, None) : literal)) (tree |> member "body" |> to_list_or_empty);
+    otherwise =
+      match tree |> member "else" with
+      | `Null -> None
+      | else_branch -> Some (parse_rule else_branch)
   }
 
+(* --- *)
+
+let parse_import_path tree =
+  match (tree |> member "type" |> to_string) with
+  | "ref" -> tree |> member "value" |> to_list |> parse_ref
+  | _ -> failwith "unimplemented"
+
+let parse_import tree =
+  { alias = tree |> member "alias" |> to_string_option;
+    path = tree |> member "path" |> parse_import_path
+  }
+
+(* --- *)
+
 let parse_package tree =
   tree |> member "path" |> to_list |> parse_ref
 
+(* --- *)
+
 let parse_tree tree =
   { path = tree |> member "package" |> parse_package;
     imports = List.map ~f:parse_import (tree |> member "imports" |> to_list_or_empty);