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

open Ast
open Core
open Lean

type variable = { lean_id : string;
                  complete : 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 });
    ("equal", { lean_id = "equal"; complete = true })
  ]

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

let extract_variables (env : environ) (q : query) : (string list * string list) =
  let combine = fun (x, y) (x', y') -> (x @ x', y @ y') in
  let rec handle_term env : term -> (string list * string list) = function
    | Call ((Var "assign", []), [ Var name; rhs ]) -> let (x, y) = handle_term env rhs in (x, [name] @ y)
    | Var name -> if Map.find env name |> Option.is_some then ([], []) else ([name], [])
    | Ref (target, args) ->
      let terms = [term_of_ref_target target] @ args in
      List.map terms ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
    | Array xs | Set xs -> List.map xs ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
    | Object kvs -> List.map ~f:snd kvs
                    |> List.map ~f:(handle_term env)
                    |> List.fold ~init:([], []) ~f:combine
    | Call (x, xs) ->
      List.map ((Ref x : term)::xs) ~f:(handle_term env) |> List.fold ~init:([], []) ~f:combine
    | ArrayComprehension _ -> failwith "unimplemented"
    | ObjectComprehension _ -> failwith "unimplemented"
    | SetComprehension _ -> failwith "unimplemented"
    | _ -> ([], [])
  and handle_expr env = function
    | Term t -> handle_term env t
    | Every (Some _key, _value, _, _q) -> failwith "unimplemented"
    | Every (None, value, _, q) ->
      let env' = Map.set env ~key:value ~data:({ lean_id = "test"; complete = true }) in
      List.map q ~f:(handle_literal env') |> List.fold ~init:([], []) ~f:combine
  and handle_literal env : literal -> (string list * string list) = function
    | Exists (exprs, _) ->
      List.map exprs ~f:(handle_expr env) |> List.fold ~init:([], []) ~f:combine
    | Expr (expr, _) -> handle_expr env expr
    | Negate (expr, _) -> handle_expr env expr
  in
  let (free, bound) = List.map q ~f:(handle_literal env) |> List.fold ~init:([], []) ~f:combine in
  let free = List.filter free  ~f:(fun x -> List.mem bound ~equal:String.equal x |> not) in
  (free |> uniq, bound |> uniq)

let extract_variables_wrapper (q : query) : (string list * string list) =
  extract_variables toplevel q

(* --- *)

type compile_term_result =
  { form : lean_expression;
    toplevel : lean_expression list;
    new_locals : (string * string) list }

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 -> compile_term_result =
  let result_of expr =
    { form = expr;
      toplevel = [];
      new_locals = []; } in
  let form = function | { form ; _ } -> form in
  let toplevel = function | { toplevel ; _ } -> toplevel in
  let new_locals = function | { new_locals ; _ } -> new_locals in

  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", []))
     |> result_of
  | Var id ->
     (match Map.find env id with
      | Some { lean_id; complete = true } -> Identifier lean_id |> result_of
      | Some { lean_id; complete = false } ->
        let anonymous_id = gensym "anonymous" in
        { form       = Identifier anonymous_id;
          toplevel   = [];
          new_locals = [(anonymous_id, lean_id)]; }
      | None ->
        failwith (Printf.sprintf "unbound variable %s" id))
  | Ref (target, terms) ->
    List.fold
      terms
      ~init:(compile_term env (term_of_ref_target target))
      ~f:(fun x y ->
          match x with
          | { form = x'; toplevel = toplevel'; new_locals = new_locals' } ->
            match compile_term env y with
            | { form = y'; toplevel = toplevel''; new_locals = new_locals'' } ->
              let toplevel = toplevel' @ toplevel'' in
              let new_locals = new_locals' @ new_locals'' in
              { form = (Application ((Identifier "Rego.Any.deref"), [x'; y']));
                toplevel; new_locals })
  | Array xs ->
    let subterms = List.map ~f:(compile_term env) xs in
    { form = Application (Identifier "Rego.Any.array", List.map ~f:form subterms);
      toplevel = List.map ~f:toplevel subterms |> List.concat;
      new_locals = List.map ~f:new_locals subterms |> List.concat }
  | Set xs ->
    let subterms = List.map ~f:(compile_term env) xs in
    { form = Application  (Identifier "Rego.Any.set", List.map ~f:form subterms);
      toplevel = List.map ~f:toplevel subterms |> List.concat;
      new_locals = List.map ~f:new_locals subterms |> List.concat }
  | Object kvs ->
     let compile_kv = function
       | (key, value) ->
         match compile_term env (term_of_object_key key) with
         | { form = x'; toplevel = toplevel'; new_locals = new_locals' } ->
           match compile_term env value with
           | { form = y'; toplevel = toplevel''; new_locals = new_locals'' } ->
             let toplevel = toplevel' @ toplevel'' in
             let new_locals = new_locals' @ new_locals'' in
             { form = Tuple [x'; y']; toplevel; new_locals } in
     let subterms = List.map ~f:compile_kv kvs in
     { form = Application  (Identifier "Rego.Any.object", List.map ~f:form subterms);
       toplevel = List.map ~f:toplevel subterms |> List.concat;
       new_locals = List.map ~f:new_locals subterms |> List.concat }
  | ArrayComprehension (_term, _query) ->
     failwith "unimplemented"
  | SetComprehension (_term, _query) ->
     failwith "unimplemented"
  | ObjectComprehension ((_key, _value), _query) ->
     failwith "unimplemented"
  | Call (fn, args) ->
    let head = compile_term env (Ref fn) in
    let subterms = List.map ~f:(compile_term env) args in
    { form = Application (form head, List.map ~f:form subterms);
      toplevel = toplevel head @ (List.map ~f:toplevel subterms |> List.concat);
      new_locals = new_locals head @ (List.map ~f:new_locals subterms |> List.concat) }

let rec compile_expr env : expr -> (environ * compile_term_result) = function
  | Term t ->
    (match t with
     (* Special case of call; this is 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 } in
       (env', compile_term env' (Call ((Var "eq", []), [ Var name; rhs ])))
     | t -> (env, compile_term env t))
  | Every (Some key, value, domain, query) ->
    (match compile_term env domain with
     | { form = domain'; toplevel; new_locals } ->
       let (toplevel', query_compiled) = compile_query env query in
       let form = ForAll
           (([key; value], "Rego.Any"),
            (Application
               (Identifier "And",
                [(Application
                    ((Identifier "Membership.mem"),
                     [Tuple [Identifier key; Identifier value]; domain']));
                 query_compiled])))
       in (env, {form; toplevel = toplevel @ toplevel'; new_locals}))
  | Every (None, value, domain, query) ->
    (match compile_term env domain with
     | { form = domain'; toplevel; new_locals } ->
       let (toplevel', query_compiled) = compile_query env query in
       let form = ForAll
           (([value], "Rego.Any"),
            (Application
               (Identifier "And",
                [(Application
                    ((Identifier "Membership.mem"),
                     [Identifier value; domain']));
                 query_compiled])))
       in (env, {form; toplevel = toplevel @ toplevel'; new_locals}))

and compile_literal env : literal -> (environ * compile_term_result) = function
  | Exists (exprs, _) ->
    let (toplevel, form) = compile_query env (List.map ~f:literal_of_expr exprs) in
    (env, { form; toplevel; new_locals = [] })
  | Expr (expr, _) -> compile_expr env expr
  | Negate (_expr, _) -> failwith "unimplemented"

and compile_query env literals : (lean_expression list * lean_expression) =
  let (free, bound) = extract_variables env literals in
  let env' = List.map (free @ bound) ~f:(fun name -> (name, { lean_id = name; complete = 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, toplevel, new_locals) =
    List.fold subterms
      ~init:([], [], [])
      ~f:(fun (a, b, c) y ->
          match y with
          | { form; toplevel; new_locals } ->
            let comparison =
              Application (Identifier "Ne", [form; Application (Identifier "Rego.Any.bool", [Bool false])])
            in
            (a @ [comparison], b @ toplevel, c @ new_locals))
  in
  (toplevel,
   if List.is_empty new_locals then
     match compiled_terms with
     | [x] -> x
     | xs -> Application (Identifier "And", xs)
   else
     Exists ((free @ bound @ (List.map ~f:fst new_locals), "Rego.Any"),
             Application (Identifier "And",
                          List.map ~f:(fun (x, y) ->
                              Application (Identifier y, [Identifier x]))
                            new_locals
                          @ 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; _ } ->
             (match compile_term toplevel term with
              (* TODO: We need to pass this up.. not ignore it... *)
              | { form; _ } ->
                (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") "₉" *)