path: root/rego-proofs/compiler-rewrite/lib/compiler.ml

open Ast
open Core
open Lean

type definition_kind = Complete
                     | Parameter
                     | ScalarRule
                     | SetRule
                     | ObjectRule
                     | FunctionRule
[@@deriving eq, show]

let kind_and_args_of_rule = function
  | { head = ((_, []), Assignment (_, term)); _ } ->
    `ScalarRule term
  | { head = ((_, _), Set term); _ } ->
    `SetRule term
  | { head = ((_, [key]), Assignment (_, value)); _ } ->
    `ObjectRule (key, value)
  | { head = ((_, _), Function (_, args, term)); _ } ->
    `FunctionRule (args, term)
  | { head = ((_, _), Assignment (_, _)); _ } as rule ->
    failwith (Printf.sprintf "unknown rule form: %s" (show_rule rule))

let definition_kind_of_rule rule =
  match kind_and_args_of_rule rule with
  | `ScalarRule _ -> ScalarRule
  | `SetRule _ -> SetRule
  | `ObjectRule _ -> ObjectRule
  | `FunctionRule _ -> FunctionRule

type variable = { lean_id : string;
                  kind : definition_kind;
                  assigned : bool;
                }
[@@deriving eq, show]

module Environ = Map.Make (String)
type environ = variable Environ.t

let gensym : string -> string =
  let is_alphanum = function  'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' -> true | _ -> false in
  let counter = ref 0 in
  fun (base : string) ->
    let base = String.filter base ~f:is_alphanum in
    let number = !counter in
    counter := !counter + 1 ;
    Printf.sprintf "v%s__%d" base number

let extend_environment (env : environ) (vars : (string * variable) list) =
  Environ.of_alist_exn vars
  |> Map.merge env ~f:(fun ~key:_key -> function
      | `Right v -> Some v
      | `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
  | Var v    -> Var v

let term_of_ref_target : ref_target -> term = function
  | Var s                        -> Var s
  | Array xs                     -> Array xs
  | Object kvs                   -> Object kvs
  | Set xs                       -> Set xs
  | ArrayComprehension (key, q)  -> ArrayComprehension (key, q)
  | ObjectComprehension (key, q) -> ObjectComprehension (key, q)
  | SetComprehension (key, q)    -> SetComprehension (key, q)
  | Call (fn, args)              -> Call (fn, args)

let literal_of_expr exp : literal = Expr (exp, None)

let map_with_environment_updates env xs fn =
  List.fold xs ~init:(env, []) ~f:(fun (env, ys) x ->
      let env', y = fn env x in
      (env', ys @ [ y ]))

let rec uniq eq = function
  | [] -> []
  | hd::tl ->
    if List.mem tl hd ~equal:eq then
      uniq eq tl
    else hd :: uniq eq tl

type analysis_result =
  { bound : string list;
    free  : string list;
    used  : string list }

let extract_variables (env : environ) (q : query) : analysis_result =
  let empty = { bound = []; free = []; used = [] } in
  let combine =
    fun x y ->
      match (x, y) with
      | ({ bound = bound';  free = free';  used = used' },
         { bound = bound''; free = free''; used = used'' }) ->
        { bound = bound' @ bound'';
          free = free' @ free'';
          used = used' @ used'' } in
  let rec handle_term env : term -> analysis_result = function
    (* If we see a variable in an assignment, and it isn't a parameter, it's bound. *)
    | Call ((Var "assign", []), [ Var name; rhs ]) ->
      let { bound; free; used } = handle_term env rhs in
      (match Map.find env name with
       | Some { kind = Parameter; _ } ->
         { used = [name] @ used; free; bound }
       | None ->
         { bound = [name] @ bound; free; used }
       | _ ->
         assert false (* TODO *))
    (* Variables that haven't previously been introduced to the scope are free. *)
    | Var name ->
      if Map.find env name |> Option.is_some then
        { free = []; bound = []; used = [name] }
      else { free = [name]; bound = []; used = [] }
    (* Otherwise, we'll recurse on the term. *)
    | Ref (target, args) ->
      let target = term_of_ref_target target in
      if List.is_empty args then
        handle_term env target
      else
        let terms = [target] @ args in
        List.map terms ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
    | Array xs | Set xs ->
      List.map xs ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
    | Object kvs -> List.map ~f:snd kvs
                    |> List.map ~f:(handle_term env)
                    |> List.fold ~init:empty ~f:combine
    | Call (x, xs) ->
      List.map ((Ref x : term)::xs) ~f:(handle_term env) |> List.fold ~init:empty ~f:combine
    | ArrayComprehension _ -> assert false (* TODO *)
    | ObjectComprehension _ -> assert false (* TODO *)
    | SetComprehension _ -> assert false (* TODO *)
    | _ -> empty
  and handle_expr env = function
    | Term t -> handle_term env t
    | Every (Some _key, _value, _, _q) -> assert false (* TODO *)
    | Every (None, value, _, q) ->
      let var = ({ lean_id = "unimportant"; kind = Complete; assigned = true }) in
      let env' = Map.set env ~key:value ~data:var in
      List.map q ~f:(handle_literal env') |> List.fold ~init:empty ~f:combine
  and handle_literal env : literal -> analysis_result = function
    | Exists (exprs, _) ->
      List.map exprs ~f:(handle_expr env) |> List.fold ~init:empty ~f:combine
    | Expr (expr, _) -> handle_expr env expr
    | Negate (expr, _) -> handle_expr env expr
  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
  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

(* --- *)

(* 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 ->
    [rule]
  | { head = ((tgt, xs), rhs); _ } as rule  ->
    let tail = List.last_exn xs in
    let name = gensym "anonymous" in
    [{ rule with head = ((Var name, [tail]), rhs) }] @
    transform_rule_with_ref_head { rule with head = ((tgt, List.drop_last_exn xs), Assignment (Assign, Var name)) }

(* --- *)

let rec compile_term env : term -> lean_expression =
  function
  | Scalar s ->
     (match s with
      | String s -> Application (Identifier "Rego.Any.str",  [String s])
      | Number x -> Application (Identifier "Rego.Any.num",  [Number x])
      | Bool b   -> Application (Identifier "Rego.Any.bool", [Bool b])
      | Null     -> Application (Identifier "Rego.Any.null", []))
  | Var id ->
    (match Map.find env id with
     | Some { lean_id; kind = Complete; assigned = true } |
       Some { lean_id; kind = Parameter; assigned = true }
       -> Identifier lean_id
     | Some { assigned = false; _ } ->
       failwith (Printf.sprintf "identifier %s used before definition" id)
     | Some var ->
       failwith (Printf.sprintf
                   "variable found in environment but of unknown form %s"
                   (show_variable var))
     | None -> failwith (Printf.sprintf "unbound variable %s" id))
  | Ref (target, terms) ->
    let target' = compile_term env (term_of_ref_target target) in
    let terms' = List.map ~f:(compile_term env) terms in
    Application (Identifier "Rego.Any.deref", [target'] @ terms')
  | Array xs ->
    let xs' = List.map ~f:(compile_term env) xs in
    Application (Identifier "Rego.Any.array", xs')
  | Set xs ->
    let xs' = List.map ~f:(compile_term env) xs in
    Application (Identifier "Rego.Any.set", xs')
  | Object kvs ->
     let compile_kv = function
       | (key, value) ->
         let key' = compile_term env (term_of_object_key key) in
         let value' = compile_term env value in
         Tuple [key'; value'] in
     let kvs' = List.map ~f:compile_kv kvs in
     Application (Identifier "Rego.Any.object", kvs')
  | ArrayComprehension (_term, _query) ->
     assert false (* TODO *)
  | SetComprehension (_term, _query) ->
     assert false (* TODO *)
  | ObjectComprehension ((_key, _value), _query) ->
     assert false (* TODO *)
  | Call (fn, args) ->
    let fn' = compile_term env (Ref fn) in
    let args' = List.map ~f:(compile_term env) args in
    Application (fn', args')

let rec compile_expr env : expr -> (environ * lean_expression) = function
  | Term t ->
    (match t with
     (* Call is a special case, being the only kind of "term" that affects scope. *)
     | Call ((Var "assign", []), [ Var name; rhs ]) ->
       (* Assumption: `name' is in the `env' because of our up-front variable analysis, but we've marked it as unassigned. *)
       let { lean_id; kind; _ } = Map.find_exn env name in
       (* assert (not assigned) ; *)

       let variable = { assigned = true; lean_id; kind } in
       let env' = Map.set env ~key:name ~data:variable in
       (
         env',
         compile_term env' (Call ((Var "eq", []), [ Var name; rhs ]))
       )
     | t -> (env, compile_term env t))
  | Every (Some key, value, domain, query) ->
    let domain' = compile_term env domain in
    (* Extend the environment with the key and value. *)
    let env' = Map.set env ~key:key ~data:{ lean_id = gensym key; kind = Complete; assigned = true }
               |> Map.set ~key:value ~data:{ lean_id = gensym value; kind = Complete; assigned = true } in
    let (env'', query') = compile_query env' query in
    (env'',
     ForAll
       (([key; value], "Rego.Any"),
        (Application
           (Identifier "And",
            [(Application
                ((Identifier "Membership.mem"),
                 [Tuple [Identifier key; Identifier value]; domain']));
             query']))))
  | Every (None, value, domain, query) ->
    let domain' = compile_term env domain in
    (* Extend the environment with the value. *)
    let env' = Map.set env ~key:value ~data:{ lean_id = gensym value; kind = Complete; assigned = true } in
    let (env'', query') = compile_query env' query in
    (env'',
     ForAll
       (([value], "Rego.Any"),
        (Application
           (Identifier "And",
            [(Application
                ((Identifier "Membership.mem"),
                 [Identifier value; domain']));
             query']))))

and compile_literal env : literal -> (environ * lean_expression) = function
  | Exists (exprs, _) ->
    let (env', form) = compile_query env (List.map ~f:literal_of_expr exprs) in
    (env', form)
  | Expr (expr, _) ->
    let (env', inner) = compile_expr env expr in
    (env', Application (Identifier "Ne", [inner; Application (Identifier "Rego.Any.bool", [Bool false])]))
  | Negate (expr, _) ->
    let (env', form) = compile_literal env (Expr (expr, None)) in
    (env', Application (Identifier "Not", [form]))

and compile_query env literals : (environ * lean_expression) =
  let { free; bound; used } = extract_variables env literals in
  let used = used
             |> List.filter ~f:(fun name -> match Map.find_exn env name with
                                            | { kind = Complete; _ } -> false
                                            | _ -> true) in
  let env' =
    List.concat
      [List.map free  ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = true }));
       List.map used  ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = true }));
       List.map bound ~f:(fun name -> (name, { lean_id = gensym name; kind = Complete; assigned = false }))]
    |> extend_environment env in
  let (env'', subterms) = map_with_environment_updates env' literals compile_literal in

  let lean_id_of_identifier k = (let { lean_id; _ } = Map.find_exn env'' k in lean_id) in
  let constraints =
    List.filter_map used ~f:(fun name ->
        match Map.find env name with
        | Some { lean_id; kind = ScalarRule; _ } ->
          Some
            (Application (Identifier "scalar_specified_by",
                          [Identifier lean_id;
                           Identifier (lean_id_of_identifier name)]))
        | Some { lean_id; kind = SetRule; _ } ->
          Some
            (Application (Identifier "set_specified_by",
                          [Identifier lean_id;
                           Identifier (lean_id_of_identifier name)]))
        | Some { lean_id; kind = ObjectRule; _ } ->
          Some
            (Application (Identifier "object_specified_by",
                          [Identifier lean_id;
                           Identifier (lean_id_of_identifier name)]))
        | Some { lean_id; kind = FunctionRule; _ } ->
          Some
            (Application (Identifier "function_specified_by",
                          [Identifier lean_id;
                           Identifier (lean_id_of_identifier name)]))
        | Some _ -> None
        | None -> failwith (Printf.sprintf "unbound variable %s" name)) in
  let compiled_terms = constraints @ subterms in

  let form = match compiled_terms with
     | [x] -> x
     | xs -> Application (Identifier "And", xs)
  in

  let quantified_names = free @ bound @ used in
  (env'',
   if List.is_empty quantified_names then
     form
   else
     Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
             form))

let compile_definition env (rules : rule list) =
  assert (not (List.is_empty rules));

  (* 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.hd_exn rules with
    | { head = (rf, _); _ } -> rf
  in

  let (default_rules, conditional_rules) =
    List.partition_tf rules ~f:(function
        | { default = true; _ } -> true
        | _ -> false)
  in
  (* There can be only one default rule. *)
  assert (List.length default_rules <= 1);

  let subjects =
    let rule = List.nth_exn rules 0 in
    let identifier_pair_of_term default_name : term -> _ = function
      | Var identifier when Map.find env identifier |> Option.is_none ->
        let anon_name = gensym identifier in
        (identifier, anon_name, Parameter, true, None)
      | term ->
        let anon_name = gensym default_name in
        let assertion = (Call ((Var "eq", []), [Var anon_name; term]) : term) in
        (anon_name, anon_name, Complete, true, Some assertion)
    in
    match rule |> kind_and_args_of_rule with
    | `ScalarRule term | `SetRule term ->
      [term |> identifier_pair_of_term "x"]
    | `ObjectRule (key, value) ->
      [key |> identifier_pair_of_term "k";
       value |> identifier_pair_of_term "v"]
    | `FunctionRule (args, output) ->
      List.map args ~f:(fun arg_name -> (arg_name, gensym arg_name, Parameter, false, None))
      @ [output |> identifier_pair_of_term "y"]
  in
  let env' =
    List.map subjects
      ~f:(fun (subject, lean_id, kind, output, _) ->
          match kind with
          | Parameter -> (subject, { lean_id = lean_id; assigned = not output; kind })
          | _ -> (subject, { lean_id = lean_id; assigned = true; kind }))
    |> extend_environment env
  in
  let conditional_arms =
    List.mapi conditional_rules ~f:(fun i rule ->
        let { body; _ } = rule in
        let (_, body_form) = compile_query env' body in
        let subject_constraints =
          List.filter_map subjects ~f:(fun (_, _, _, _, x) -> x)
          |> List.map ~f:(fun clause ->
              let (_, compiled) = [Expr (Term clause, None)] |> compile_query env' in
              compiled) in
        (body_form,
         (Printf.sprintf "rule%d" i,
          subject_constraints @ [body_form],
          List.map subjects ~f:(fun (_, x, _, _, _) -> Identifier x)))) in
  let default_arms =
    List.map default_rules ~f:(fun _ ->
        let other_rule_bodies = List.map ~f:fst conditional_arms in
        let subject_constraints =
          List.filter_map subjects ~f:(fun (_, _, _, _, x) -> x)
          |> List.map ~f:(fun clause ->
              let (_, compiled) = [Expr (Term clause, None)] |> compile_query env' in
              compiled) 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_constraints @ [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
      | _ -> assert false (* unreachable *)
    in
    Inductive
      (string_of_reference variable_being_defined,
       [(List.map subjects ~f:(fun (_, x, _, _, _) -> x), "Rego.Any")],
       List.append (List.map ~f:snd conditional_arms) default_arms)
  in (env, form)

let compile_policy env (rules : rule list) =
  let group accessor l =
    let rec grouping acc = function
      | [] -> acc
      | hd::tl ->
        let l1,l2 = List.partition_tf ~f:(fun x -> equal_reference (accessor x) (accessor hd)) tl in
        grouping ((hd::l1)::acc) l2
    in
    grouping [] l in
  let head_of_rule = function | {head = (rf, _); _ } -> rf in
  let by_reference = group head_of_rule rules in
  let env' =
    List.map by_reference ~f:(function
        | { head = ((Var name, []), Assignment _); _ } :: _ ->
          (name, { lean_id = name ; kind = ScalarRule; assigned = true })
        | { head = ((Var name, [_]), Assignment _); _ } :: _ ->
          (name, { lean_id = name ; kind = ObjectRule; assigned = true })
        | { head = ((Var name, []), Set _); _ } :: _ ->
          (name, { lean_id = name ; kind = SetRule; assigned = true })
        | { head = ((Var name, []), Function _); _ } :: _ ->
          (name, { lean_id = name ; kind = FunctionRule; assigned = true })
        | _ -> failwith "policy rule head is not a variable")
    |> extend_environment env in
  map_with_environment_updates env' by_reference compile_definition

let compile_package = function
  | { path; imports; policy } ->
    let policy = List.map ~f:transform_rule_with_ref_head policy |> List.concat in
    let (_env', forms) = compile_policy toplevel policy in
    [Comment (show_reference path |> String.escaped);
     Comment (List.map ~f:show_import imports |> String.concat |> String.escaped)]
    @ forms