path: root/rego-proofs/compiler-rewrite/bin/main.ml

open Core
open Yojson.Basic.Util

type expr = Term of term
          | Every of (string option) * string * term * query
[@@deriving show]
and literal = Exists of (expr list) * (with_modifier option)
            | Expr of expr * (with_modifier option)
            | Negate of expr * (with_modifier option)
[@@deriving show]
and query = literal list
[@@deriving show]
and term = Scalar of scalar
          | Var of string
          | Ref of reference
          | Array of term list
          | Object of (object_key * term) list
          | Set of term list
          | ArrayComprehension of term * query
          | ObjectComprehension of (object_key * term) * query
          | SetComprehension of term * query
          | Call of reference * (term list)
[@@deriving show]
and scalar = String of string
           | Number of float
           | Bool of bool
           | Null
[@@deriving show]
and reference = ref_target * (term list)
[@@deriving show]
and ref_target = Var of string
          | Array of term list
          | Object of (object_key * term) list
          | Set of term list
          | ArrayComprehension of term * query
          | ObjectComprehension of (object_key * term) * query
          | SetComprehension of term * query
          | Call of reference * (term list)
[@@deriving show]
and object_key = Scalar of scalar
               | Ref of reference
               | Var of string
[@@deriving show]
and with_modifier = (term * term) list
[@@deriving show]

type assign_operator = Assign | Unify
[@@deriving show]

type rule_head_lhs = reference
[@@deriving show]
type rule_head_rhs = Assignment of assign_operator * term
                   | Function of assign_operator * (string list) * term
                   | Set of term
[@@deriving show]
type rule_head = rule_head_lhs * rule_head_rhs
[@@deriving show]
type rule =
  { default : bool;
    head : rule_head;
    body: query;
    otherwise : rule option;
  }
[@@deriving fields, show]

type import =
  { alias : string option;
    path : reference;
  }
[@@deriving show]

type package =
  { path : reference;
    imports: import list;
    policy : rule list;
  }
[@@deriving fields, show]

module Environ = Map.Make(String);;

let extract_local_variables (q : query) (scope : string Environ.t) : (string list) =
  let handle_term = function
    | Call (((Var "assign"), []), [(Var name); _]) ->
       if (Map.find scope name) |> Option.is_none then [name] else []
    | _ -> []
    (* | Var name -> if (Map.find scope name) |> Option.is_none then [name] else [] *)
    (* | Array of term list *)
    (* | Object of (object_key * term) list *)
    (* | Set of term list *)
  in
  let handle_expr = function
    | Term t -> handle_term t
    | Every ((Some _key), _value, _, _q) -> failwith "unimplemented"
    | Every (None, _value, _, _q) -> failwith "unimplemented"
  in
  let handle_literal = function
    | Exists (exprs, _) -> List.map ~f:handle_expr exprs |> List.concat
    | Expr (expr, _) -> handle_expr expr
    | Negate (expr, _) -> handle_expr expr
  in
  List.map q ~f:handle_literal |> List.concat

(* --- *)

let is_null : Yojson.Basic.t -> bool = function
  | `Null -> true
  | _ -> false

let to_bool_or_false : Yojson.Basic.t -> bool = function
  | `Bool b -> b
  | _ -> false

let to_list_or_empty : Yojson.Basic.t -> Yojson.Basic.t list = function
  | `Null -> []
  | `List xs -> xs
  | _ -> failwith "Not a list or null"

let object_key_of_term : term -> object_key = function
  | Scalar s -> Scalar s
  | Ref r -> Ref r
  | Var v -> Var v
  | _ -> failwith "failed to convert term to object_key"

let rec parse_term (tree : Yojson.Basic.t) : term =
  let term_type = tree |> member "type" |> to_string in
  match term_type with
  | "var" -> Var (tree |> member "value" |> to_string)
  | "ref" -> Ref (tree |> member "value" |> to_list |> parse_ref)
  | "boolean" -> Scalar (Bool (tree |> member "value" |> to_bool))
  | "number" -> Scalar (Number (tree |> member "value" |> to_number))
  | "string" -> Scalar (String (tree |> member "value" |> to_string))
  | "array" -> Array (tree |> member "value" |> to_list |> List.map ~f:parse_term)
  | "set" -> Set (tree |> member "value" |> to_list |> List.map ~f:parse_term)
  | "object" ->
     Object
       (tree
        |> member "value"
        |> to_list
        |> List.map ~f:(function
               | `List [x; y] ->
                  (x |> parse_term |> object_key_of_term, y |> parse_term)
               | _ -> failwith "ill-formed object"))
  | "arraycomprehension" ->
     ArrayComprehension
       (tree |> member "value" |> member "term" |> parse_term,
        tree |> member "value" |> member "body" |> to_list |> List.map ~f:parse_literal)
  | "setcomprehension" ->
     SetComprehension
       (tree |> member "value" |> member "term" |> parse_term,
        tree |> member "value" |> member "body" |> to_list |> List.map ~f:parse_literal)
  | "objectcomprehension" ->
     ObjectComprehension
       ((tree |> member "value" |> member "key" |> parse_term |> object_key_of_term,
         tree |> member "value" |> member "value" |> parse_term),
        tree |> member "value" |> member "body" |> to_list |> List.map ~f:parse_literal)
  | "call" ->
     (match tree |> member "value" |> to_list with
      | func :: args ->
         Call (func |> member "value" |> to_list |> parse_ref,
               List.map ~f:parse_term args)
      | _ -> failwith "ill-formed call")
  | _ -> failwith ("unsupported term type " ^ term_type)

and parse_literal tree : literal =
  let with_expr = tree |> member "with" in
  let with_expr = match with_expr with
    | `Null -> None
    | `List xs ->
       Some
         (List.map xs ~f:(fun x -> (x |> member "target" |> parse_term,
                                   x |> member "value" |> parse_term)))
    | _ -> failwith ("ill-formed with expression " ^ (Yojson.Basic.to_string with_expr))
  in
  match tree |> member "terms" with
  | `List (func :: args) ->
     Expr (Term (Call (func |> member "value" |> to_list |> parse_ref,
                       List.map ~f:parse_term args)),
           with_expr)
  | `List _ -> failwith "ill-formed call"
  | x ->
     if not (x |> member "symbols" |> is_null) then
       Exists (x |> member "symbols"
                 |> to_list
                 |> List.map ~f:(fun x -> Term (parse_term x)),
               with_expr)
     else if x |> member "domain" |> is_null then
       Expr (Term (parse_term x), with_expr)
     else
       Expr (Every ((if x |> member "key" |> is_null then
                       None
                     else
                       Some (x |> member "key" |> member "value" |> to_string)),
                    x |> member "value" |> member "value" |> to_string,
                    x |> member "domain" |> parse_term,
                    (List.map ~f:parse_literal (x |> member "body" |> to_list))),
             with_expr)

and parse_ref_target tree : ref_target =
  let term = parse_term tree in
  match term with
  | Var term -> Var term
  | Array term -> Array term
  | Object term -> Object term
  | Set term -> Set term
  | ArrayComprehension (term, y) -> ArrayComprehension (term, y)
  | ObjectComprehension (term, y) -> ObjectComprehension (term, y)
  | SetComprehension (term, y) -> SetComprehension (term, y)
  | Call (term, y) -> Call (term, y)
  | _ -> failwith ("unsupported ref_target " ^ (show_term term))

and parse_ref tree =
  match tree with
  | target :: args -> (parse_ref_target target, List.map ~f:parse_term args)
  | _ -> failwith "empty reference"

let parse_head 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 not (tree |> member "key" |> is_null) then
      Set (tree |> member "key" |> 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 (assign_operator, (tree |> member "value" |> parse_term))
  in (lhs, rhs)

let rec parse_rule tree =
  { default = tree |> member "default" |> to_bool_or_false;
    head = tree |> member "head" |> parse_head;
    body = List.map ~f:parse_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);
    policy = List.map ~f:parse_rule (tree |> member "rules" |> to_list_or_empty)
  }

let invoke ?(stdin="") ?(env=[| |]) cmd =
  let ic, oc, ec = UnixLabels.open_process_full cmd ~env:env in
  Out_channel.output_string oc stdin ;
  Out_channel.close oc;
  (In_channel.input_all ic,
   In_channel.input_all ec)

let () =
  let file_name = Array.get (Sys.get_argv ()) 1 in
  let policy_code = In_channel.read_all file_name in
  let (stdout, _) = invoke ("/home/jakob/.local/bin/opa parse -f json /dev/stdin") ~stdin:policy_code ~env:[| |] in
  let ast = Yojson.Basic.from_string stdout |> parse_tree in
  match ast with
  | { policy = [{ body = body; _ }]; _ } -> (extract_local_variables body Environ.empty) |> List.iter ~f:print_endline
  | _ -> print_endline "failed"