Complete parsing of rule_head_rhs

1 files changed, 14 insertions, 7 deletions

diff --git a/rego-proofs/compiler-rewrite/bin/main.ml b/rego-proofs/compiler-rewrite/bin/main.ml
index 14a1ff7..11ecc44 100644
--- a/rego-proofs/compiler-rewrite/bin/main.ml
+++ b/rego-proofs/compiler-rewrite/bin/main.ml
@@ -48,9 +48,8 @@ type assign_operator = Assign | Unify
 type rule_head_lhs = reference
 [@@deriving show]
 type rule_head_rhs = Assignment of assign_operator * term
-                   | Object of term * assign_operator * term
-                   | Function of (term list) * assign_operator * term
-                   | Set of term
+                   | Function of assign_operator * (string list) * term
+                   | Set of string * term
 [@@deriving show]
 type rule_head = rule_head_lhs * rule_head_rhs
 [@@deriving show]
@@ -189,12 +188,20 @@ let parse_head tree =
   Printf.printf "%s: %s\n" "parse_head" (Yojson.Basic.to_string tree);
   (* Variables are desugared to empty references. *)
   let lhs = tree |> member "ref" |> to_list |> parse_ref in
+  let assign_operator = if tree |> member "assign" |> to_bool_or_false then
+                          Assign
+                        else
+                          Unify in
   let rhs =
-    if tree |> member "assign" |> to_bool_or_false then
-      Assignment (Assign, (tree |> member "value" |> parse_term))
-                 (* TODO: Rest cases? *)
+    if not (tree |> member "key" |> is_null) then
+      Set (tree |> member "key" |> to_string,
+           (tree |> member "value" |> parse_term))
+    else if tree |> member "args" |> to_bool_or_false then
+      Function (assign_operator,
+                (tree |> member "args" |> to_list |> List.map ~f:(fun x -> (x |> member "value" |> to_string))),
+                (tree |> member "value" |> parse_term))
     else
-      Assignment (Unify, (tree |> member "value" |> parse_term))
+      Assignment (assign_operator, (tree |> member "value" |> parse_term))
   in (lhs, rhs)
 
 let rec parse_rule tree =