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

open Ast
open Core
open Lean

type definition_kind = Complete
                     | ScalarRule
                     | SetRule
                     | ObjectRule
                     | FunctionRule

type variable = { lean_id : string;
                  kind : definition_kind;
                  assigned : bool;
                }

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

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 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 toplevel = Environ.of_alist_exn [
    ("eq",    { lean_id = "eq";    kind = Complete; assigned = true });
    ("equal", { lean_id = "equal"; kind = Complete; assigned = true })
  ]

let rec uniq = function
  | [] -> []
  | hd::tl ->
    if List.mem tl hd ~equal:String.equal then
      uniq tl
    else hd :: uniq 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, it's bound. *)
    | Call ((Var "assign", []), [ Var name; rhs ]) ->
      let { bound; free; used } = handle_term env rhs in
      { bound = [name] @ bound; free; used }
    (* 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 = [] @ 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 _ -> failwith "unimplemented"
    | ObjectComprehension _ -> failwith "unimplemented"
    | SetComprehension _ -> failwith "unimplemented"
    | _ -> empty
  and handle_expr env = function
    | Term t -> handle_term env t
    | Every (Some _key, _value, _, _q) -> failwith "unimplemented"
    | 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
  { free = free |> uniq; bound = bound |> uniq; used = used |> uniq }

let extract_variables_wrapper (q : query) : analysis_result =
  extract_variables toplevel q

(* --- *)

let gensym : string -> string =
  let counter = ref 0 in
  fun (base : string) ->
    let number = !counter in
    counter := !counter + 1 ;
    Printf.sprintf "%s__%d" base number

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 } -> Identifier lean_id
     | Some { assigned = false; _ } ->
       failwith (Printf.sprintf "identifier %s used before definition" id)
     | Some _ -> failwith "unimplemented"
     | 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) ->
     failwith "unimplemented"
  | SetComprehension (_term, _query) ->
     failwith "unimplemented"
  | ObjectComprehension ((_key, _value), _query) ->
     failwith "unimplemented"
  | 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; assigned } = 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, _) -> compile_expr env expr
  | Negate (_expr, _) -> failwith "unimplemented"

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 compiled_terms = List.map subterms
      ~f:(fun x -> Application (Identifier "Ne", [x; Application (Identifier "Rego.Any.bool", [Bool false])])) in

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

  let quantified_names = free @ bound @ used in
  (env'',
   if List.is_empty quantified_names then
     match compiled_terms with
     | [x] -> x
     | xs -> Application (Identifier "And", xs)
   else
     Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
             Application (Identifier "And", compiled_terms)))

(* TODO: Probably get rid of this once we deal with nested definitions. *)
let string_of_reference : reference -> string = function
  | (Var name, _) -> name
  | _ -> failwith "unimplemented"

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

  let variable_being_defined =
    match List.nth_exn rules 0 with
    | { head = (rf, _); _ } -> rf
  in

  let subject = gensym "x" in
  let form =
    Inductive
      (string_of_reference variable_being_defined,
       [([subject], "Rego.Any")],
       List.mapi rules ~f:(fun i ->
           function
           | { head = (_, Assignment (_, term)); body; _ } ->
              let eq : term = Call ((Var "eq", []), [Var subject; term]) in
              let env' = [(subject, { lean_id = subject; kind = Complete; assigned = true })]
                         |> extend_environment env in
              let (_, form) = compile_query env' [Expr (Term eq, None)] in
              (Printf.sprintf "rule%d" i,
               [form; compile_query env body |> snd],
               [Identifier subject])
           | _ -> failwith "unimplemented"))
  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 _scope = List.map by_reference *)
  (*     ~f:(fun x -> *)
  (*         let lean_id = List.nth_exn x 0 |> head_of_rule |> string_of_reference in *)
  (*         lean_id *)
  (*       ) *)
  (* 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 =
      List.map policy ~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 toplevel in

    let (_env', forms) = compile_policy env policy in
    [Comment (show_reference path |> String.escaped);
     Comment (List.map ~f:show_import imports |> String.concat |> String.escaped)]
    @ forms

(* let subscript_of_number n = *)
(*   string_of_int n *)
(*   |> Str.global_replace (Str.regexp "0") "₀" *)
(*   |> Str.global_replace (Str.regexp "1") "₁" *)
(*   |> Str.global_replace (Str.regexp "2") "₂" *)
(*   |> Str.global_replace (Str.regexp "3") "₃" *)
(*   |> Str.global_replace (Str.regexp "4") "₄" *)
(*   |> Str.global_replace (Str.regexp "5") "₅" *)
(*   |> Str.global_replace (Str.regexp "6") "₆" *)
(*   |> Str.global_replace (Str.regexp "7") "₇" *)
(*   |> Str.global_replace (Str.regexp "8") "₈" *)
(*   |> Str.global_replace (Str.regexp "9") "₉" *)