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

open Ast
open Core
open Lean

type variable = { lean_id : string;
                  complete : bool;
                  assigned : bool;
                }

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

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";    complete = true; assigned = true });
    ("equal", { lean_id = "equal"; complete = true; 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 terms = [term_of_ref_target 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 _ -> 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"; complete = true; 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 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; complete = true; assigned = true } -> Identifier lean_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 ]) ->
       let env' = Map.set env ~key:name ~data:{ lean_id = name; complete = true; assigned = true } 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 = key; complete = true; assigned = true }
               |> Map.set ~key:value ~data:{ lean_id = value; complete = true; 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 = value; complete = true; 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; _ } = extract_variables env literals in
  let env' = List.map free ~f:(fun name -> (name, { lean_id = name; complete = true; assigned = true }))
             |> Environ.of_alist_exn
             |> Map.merge env ~f:(fun ~key:_key -> function
                 | `Right v -> Some v
                 | `Left v -> Some v
                 | `Both (_, v) -> Some v) 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 quantified_names = free @ bound in
  (env'',
   if List.is_empty quantified_names then
     match compiled_terms with
     | [x] -> x
     | xs -> Application (Identifier "And", xs)
   else
     Exists ((free @ bound, "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 form =
    Inductive
      (string_of_reference variable_being_defined,
       [([gensym "x"], "Rego.Any")],
       List.mapi rules ~f:(fun i ->
           function
           | { head = (_, Assignment (_, term)); body; _ } ->
             let form = compile_term toplevel term in
             (Printf.sprintf "rule%d" i,
              compile_query toplevel body |> snd,
              [form])
           | _ -> 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 env = 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") "₉" *)