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

open Ast
open Core
open Lean

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

let fold_or props =
  match List.reduce props ~f:(fun x y -> Application (Identifier "Or", [y; x])) with
  | Some x -> x
  | None -> Bool true

let fold_and props =
  match List.reduce props ~f:(fun x y -> Application (Identifier "And", [y; x])) with
  | Some x -> x
  | None -> Bool false

let fold_deref terms =
  List.reduce_exn terms ~f:(fun x y -> Application (Identifier "Rego.Any.deref", [x; y]))

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 ref_target_of_term : term -> ref_target = 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)
  | _                            -> failwith "cannot widen to ref_target"

let literal_of_expr exp : literal = Expr (exp, None)
let literal_of_term t = Expr (Term t, 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 }
[@@deriving show]

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 }
       | _ ->
         { 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 = [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
    | _ -> empty
  and handle_expr env = function
    | Term t -> handle_term env t
    | Every (Some key, value, _, q) ->
      let env' = [(key,   { lean_id = "unimportant"; kind = Complete; assigned = true });
                  (value, { lean_id = "unimportant"; kind = Complete; assigned = true })]
                 |> extend_environment env in
      List.map q ~f:(handle_literal env') |> List.fold ~init:empty ~f:combine
    | 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 used = List.filter used  ~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 simplify_expression = function
  | Application (Identifier "And", xs) ->
    let xs' = List.map ~f:simplify_expression xs
              |> List.filter ~f:(function Bool true -> false | _ -> true)
    in
    (match xs' with
     | [] -> Bool false
     | [x] -> x
     | xs' -> Application (Identifier "And", xs'))
  | Application (Identifier "Eq", [Application (Identifier "Rego.Any.bool", [Bool a]);
                                   Application (Identifier "Rego.Any.bool", [Bool b])]) ->
    Bool (equal_bool a b)
  | Application (Identifier "Ne", [Application (Identifier "Rego.Any.bool", [Bool a]);
                                   Application (Identifier "Rego.Any.bool", [Bool b])]) ->
    Bool (not (equal_bool a b))
  | Application (Identifier "Eq", [Application (Identifier "Rego.Any.str", [String a]);
                                   Application (Identifier "Rego.Any.str", [String b])]) ->
    Bool (equal_string a b)
  | Application (Identifier "Ne", [Application (Identifier "Rego.Any.str", [String a]);
                                   Application (Identifier "Rego.Any.str", [String b])]) ->
    Bool (not (equal_string a b))
  | Application (Identifier "Eq", [Application (Identifier "Rego.Any.num", [Number a]);
                                   Application (Identifier "Rego.Any.num", [Number b])]) ->
    Bool (equal_float a b)
  | Application (Identifier "Ne", [Application (Identifier "Rego.Any.num", [Number a]);
                                   Application (Identifier "Rego.Any.num", [Number b])]) ->
    Bool (not (equal_float a b))
  | expr -> expr

(* --- *)

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
    if List.is_empty terms' then
      target'
    else [target'] @ terms' |> fold_deref
  | 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 "ArrayComprehension encountered, should have been lifted"
  | SetComprehension (_term, _query) ->
     failwith "ArrayComprehension encountered, should have been lifted"
  | ObjectComprehension ((_key, _value), _query) ->
     failwith "ArrayComprehension encountered, should have been lifted"
  | 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 lift_term (env : environ) (rule : rule) (f : environ -> term -> ('a * term) option) : (('a * term) list * rule) =
  let empty = ([], []) in
  let combine = fun (side_info, q) (side_info', term) ->
    (side_info @ side_info', q @ [term]) in
  let rec handle_term : term -> (('a * term) list * term) = function
    | term when f env term |> Option.is_some ->
      let (side_info, term') = f env term |> Option.value_exn in
      ([(side_info, term)], term')
    | Ref (target, args) ->
      let target = term_of_ref_target target in
      let terms = [target] @ args in

      let (side_info, terms') =
        List.map terms ~f:handle_term
        |> List.fold ~f:combine ~init:empty
      in (side_info, Ref (ref_target_of_term (List.hd_exn terms'), List.tl_exn terms'))
    | Array xs ->
      let (side_info, terms) =
        List.map xs ~f:handle_term
        |> List.fold ~f:combine ~init:empty
      in (side_info, Array terms)
    | Set xs ->
      let (side_info, terms) =
        List.map xs ~f:handle_term
        |> List.fold ~f:combine ~init:empty
      in (side_info, Set terms)
    | Object kvs ->
      let (side_info, terms) =
        List.map ~f:snd kvs
        |> List.map  ~f:handle_term
        |> List.fold ~f:combine ~init:empty
      in
      (side_info,
       Object
         (List.zip_exn
            (List.map ~f:fst kvs)
            terms))
    | Call (x, xs) ->
      let (side_info, terms) =
        List.map ((Ref x : term)::xs) ~f:handle_term
        |> List.fold ~f:combine ~init:empty
      in
      (match List.hd_exn terms with
       | Ref reference ->
         (side_info, Call (reference, List.tl_exn terms))
       | term -> failwith (show_term term |> Printf.sprintf "invalid call target %s"))
    | term -> ([], term)
  and handle_expr = function
    | Term t ->
      let (side_info, t') = handle_term t in (side_info, Term t')
    | Every (k, v, t, q) ->
      let (side_info, q') =
        List.map q ~f:handle_literal
        |> List.fold ~f:(fun (side_info, q) (side_info', term) ->
            (side_info @ side_info', q @ [term]))
          ~init:empty
      in (side_info, Every (k, v, t, q'))
  and handle_literal : literal -> (('a * term) list * literal) = function
    | Exists (exprs, with_statement) ->
      let (side_info, exprs') =
        List.map exprs ~f:handle_expr
        |> List.fold ~init:empty ~f:combine
      in (side_info, Exists (exprs', with_statement))
    | Expr (expr, with_statement) ->
      let (side_info, expr') = handle_expr expr in
      (side_info, Expr (expr', with_statement))
    | Negate (expr, with_statement) ->
      let (side_info, expr') = handle_expr expr in
      (side_info, Negate (expr', with_statement))
  in
  let { body = q; head; _ } = rule in

  let (side_info, query) = List.map ~f:handle_literal q |> List.fold ~init:empty ~f:combine in
  let (side_info', head') = match head with
    | (rf, Assignment (op, term)) ->
      let (side_info, term') = handle_term term in
      (side_info, (rf, Assignment (op, term')))
    | (rf, Function (op, args, term)) ->
      let (side_info, term') = handle_term term in
      (side_info, (rf, Function (op, args, term')))
    | (rf, Set term) ->
      let (side_info, term') = handle_term term in
      (side_info, (rf, Set term'))
      in
  (side_info @ side_info', { rule with body = query; head = head' })


let rewrite_with (env : environ) (q : query) : ((string * string) list * lean_expression list) =
  let empty = ([], []) in
  let combine (x1, y1) (x2, y2) = (x1 @ x2, y1 @ y2) in
  let rec handle_term substitutions : term -> ((string * string) list * lean_expression list) = function
    | Var name as v ->
      (match List.Assoc.find ~equal:equal_term substitutions v with
       | Some value ->
         let new_identifier = gensym name in
         ([(name, new_identifier)],
          [Application (Identifier "Eq", [Identifier new_identifier;
                                          compile_term env value])])
       | None -> empty)
    | Ref (target, args) ->
      let target = term_of_ref_target target in
      let terms = [target] @ args in
      List.map terms ~f:(handle_term substitutions) |> List.fold ~init:empty ~f:combine
    | Array xs | Set xs ->
      List.map xs ~f:(handle_term substitutions) |> List.fold ~init:empty ~f:combine
    | Object kvs -> List.map ~f:snd kvs
                    |> List.map ~f:(handle_term substitutions)
                    |> List.fold ~init:empty ~f:combine
    | Call (x, xs) ->
      List.map ((Ref x : term)::xs) ~f:(handle_term substitutions) |> List.fold ~init:empty ~f:combine
    | ArrayComprehension _ -> assert false (* TODO *)
    | ObjectComprehension _ -> assert false (* TODO *)
    | SetComprehension _ -> assert false (* TODO *)
    | _ -> empty
  and handle_expr substitutions = function
    | Term t -> handle_term substitutions t
    | Every (_, _, _, q) ->
      List.map q ~f:handle_literal |> List.fold ~init:empty ~f:combine
  and handle_literal : literal -> ((string * string) list * lean_expression list) = function
    | Exists (exprs, Some with_statement) ->
      List.map exprs ~f:(handle_expr with_statement)
        |> List.fold ~init:empty ~f:combine
    | Expr (expr, Some with_statement) ->
      handle_expr with_statement expr
    | Negate (expr, Some with_statement) ->
      handle_expr with_statement expr
    | _ ->
      empty
  in
  List.map ~f:handle_literal q |> List.fold ~init:empty ~f:combine

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) || not (equal_definition_kind kind Complete)) ;

       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"),
        fold_and
          [(Application
              ((Identifier "Membership.mem"),
               [Tuple [Identifier key; Identifier value]; domain']));
           query']
        |> simplify_expression))
  | 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"),
        fold_and
          [(Application
              ((Identifier "Membership.mem"),
               [Identifier value; domain']));
           query']
        |> simplify_expression))

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 (env', forms) = map_with_environment_updates env literals compile_literal in
  let form = match forms with
    | [x] -> x
    | xs -> xs |> fold_and |> simplify_expression in
  (env', form)

and compile_comprehension env anonymous_name (comp : term) =
  match comp with
  | ArrayComprehension (value, q) ->
    let { free; _ } = extract_variables env [literal_of_term value] in
    let env = List.map free ~f:(fun identifier ->
        (identifier, { lean_id = (gensym identifier); kind = Parameter; assigned = true }))
              |> extend_environment env in
    let (_, form) = compile_query env q in
    Application (Identifier "array_specified_by", [Identifier anonymous_name; form])
  | SetComprehension (value, q) ->
    let { free; _ } = extract_variables env [literal_of_term value] in
    let env = List.map free ~f:(fun identifier ->
        (identifier, { lean_id = (gensym identifier); kind = Parameter; assigned = true }))
              |> extend_environment env in
    let (_, form) = compile_query env q in
    Application (Identifier "set_specified_by", [Identifier anonymous_name; form])
  | ObjectComprehension ((key, value), q) ->
    let { free; _ } = extract_variables env [literal_of_term (term_of_object_key key)] in
    let { free = free'; _ } = extract_variables env [literal_of_term value] in
    let free = free @ free' in
    let env = List.map free ~f:(fun identifier ->
        (identifier, { lean_id = (gensym identifier); kind = Parameter; assigned = true }))
              |> extend_environment env in
    let (_, form) = compile_query env q in
    Application (Identifier "object_specified_by", [Identifier anonymous_name; form])
  | _ -> failwith (show_term comp |>
                   (Printf.sprintf "called compile_comprehension on non-comprehension %s"))

and compile_rule env subjects rule : (environ * lean_expression) list =
  let (comprehensions, rule) = lift_term env rule
      (fun _env term ->
         let anonymous_name = gensym "comprehension" in
         match term with
         | ArrayComprehension _ | SetComprehension _ | ObjectComprehension _ ->
           Some (anonymous_name, Var anonymous_name)
         | _ -> None) in
  (* Each comprehension is now mapped to an anonymous name. We'll add it to the
     scope so we don't have an unbound variable error when we compile the
     modified rule. *)
  let env = List.map comprehensions
      ~f:(fun (anonymous_name, _) ->
          (anonymous_name, { lean_id = anonymous_name; kind = Parameter; assigned = true }))
            |> extend_environment env in

  let { body = literals; _ } = rule in

  (* TODO: Deal with user-defined functions. *)
  (* let (_, literals) = lift_term env literals *)
  (*     (fun env term -> *)
  (*        match term with *)
  (*        | Call ((Var identifier, []), _args) -> *)
  (*          (match Map.find env identifier with *)
  (*           | Some { kind = FunctionRule ; _ } -> *)
  (*             let anon_name = gensym "function_call" in *)
  (*             Some (anon_name, Var anon_name) *)
  (*           | _ -> None) *)
  (*        | _ -> None) in *)


  let (new_identifiers, _with_statement_constraints) = rewrite_with env literals in
  let env = List.map new_identifiers
      ~f:(fun (old_name, anonymous_name) ->
          (old_name, { lean_id = anonymous_name; kind = Parameter; assigned = true }))
            |> extend_environment env in

  let { free; bound; used } = extract_variables env literals in
  let (used', parameters, env_updates, subject_constraints) =
    let explicit_parameters = match rule |> kind_and_args_of_rule with
      | `FunctionRule (args, _) ->
        List.filter_map args ~f:(function
            | Var id -> Some id
            | _ -> None)
      | _ -> []
    in
    let candidates = match rule |> kind_and_args_of_rule with
      | `ScalarRule term | `SetRule term ->
        [term]
      | `ObjectRule (key, value) ->
        [key; value]
      | `FunctionRule (args, output) ->
        args @ [output]
    in

    let env_updates_and_constraints =
      List.map2_exn
        subjects
        candidates
        ~f:(fun subject ->
            function
            (* In this case, we're dealing with a formal parameter, so we want
                to have it be an alias of the lean_id for the subject. *)
            | Var identifier when List.mem (bound @ explicit_parameters) identifier ~equal:String.equal ->
              (Some identifier,
               [(identifier, { lean_id = subject; kind = Parameter; assigned = true })],
               None)
            | term ->
              let { free; _ } = extract_variables env [literal_of_term term] in
              (None,
               [(subject, { lean_id = subject; kind = Parameter; assigned = true })]
               @ List.map free
                 ~f:(fun name ->
                     (name, { lean_id = gensym "param"; kind = Parameter; assigned = true })),
               Some (Call ((Var "eq", []), [Var subject; term]) : term)))
    in

    let parameters = env_updates_and_constraints |> List.filter_map ~f:(fun (x, _, _) -> x) in
    let non_parameters = List.filter_map candidates ~f:(function
        | Var id -> Some id
        | _ -> None) in

    (non_parameters,
     parameters,
     env_updates_and_constraints |> List.map ~f:(fun (_, x, _) -> x) |> List.concat,
     env_updates_and_constraints |> List.filter_map  ~f:(fun (_, _, x) -> x))
  in
  let env' = extend_environment env env_updates in
  let bound = List.filter bound ~f:(fun identifier -> not (List.mem parameters identifier ~equal:String.equal)) in

  let free = free |>
             List.filter ~f:(fun name -> not (List.mem parameters name ~equal:String.equal)) in

  let used = used @ used' in
  let used = used |>
             List.filter ~f:(fun name -> not (List.mem parameters name ~equal:String.equal)) 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 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 subject_constraints = List.map subject_constraints ~f:literal_of_term in
  let (env'', form) = compile_query env' (literals @ subject_constraints) in
  let comprehension_constraints =
    List.map comprehensions ~f:(fun (name, form) -> compile_comprehension env name form) in
  let quantified_names = free @ bound @ used in

  (env'',
   if List.is_empty quantified_names then
     fold_and [form] |> simplify_expression
   else
     Exists ((List.map quantified_names ~f:lean_id_of_identifier , "Rego.Any"),
             constraints @ comprehension_constraints @ [form]
             |> fold_and
             |> simplify_expression))
  :: (match rule.otherwise with
      | Some branch -> compile_rule env subjects branch
      | None -> [])

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 =
    match List.hd_exn rules |> kind_and_args_of_rule with
    | `ScalarRule _ | `SetRule _ ->
       [gensym "x"]
    | `ObjectRule _ ->
       [gensym "k"; gensym "v"]
    | `FunctionRule (args, _) ->
       List.map args ~f:(fun _ -> gensym "arg")
       @ [gensym "y"]
  in
  let conditional_rules =
    List.map ~f:(compile_rule env subjects) conditional_rules
    |> List.concat in
  let conditional_arms =
    List.mapi conditional_rules ~f:(fun i (_env', body) ->
        (body,
         (Printf.sprintf "rule%d" i,
          [body],
          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 body_form =
          Application (Identifier "Not", [fold_or other_rule_bodies])
        in
        ("rule_default",
         [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,
       [(subjects, "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_error_condition env =
  Map.to_alist env
  |> List.filter_map ~f:(fun (k, v) ->
      match v with
      | { kind = ScalarRule; _ } ->
        let anon_name = gensym "the" in
        Some(
          ExistsUnique (
            ([anon_name], "Rego.Any"),
            Application (Identifier "scalar_specified_by",
                         [Identifier anon_name; Identifier k])))
      | { kind = FunctionRule; _ } ->
        let args_name = gensym "args" in
        let value_name = gensym "value" in
        Some(
          ForAll
            (([args_name], "Rego.Any"),
             ExistsUnique (
               ([value_name], "Rego.Any"),
               Application
                 (Identifier "eq",
                  [Application (Identifier k, [Identifier args_name]);
                   Identifier value_name]))))
      | _ -> None)
  |> fold_and

let compile_package =
    let string_of_reference : reference -> string = function
      | (Var name, xs) ->
        [name]
        @ List.map xs ~f:(function
            | Scalar (String x) -> x
            | _ -> failwith "unreachable")
      |> String.concat ~sep:"."

      | _ -> assert false (* unreachable *)
    in
  let compile_import  = function
    | { alias = None; path } ->
      Printf.sprintf "imports %s"
        (string_of_reference path)
      |> String.escaped
    | { alias = Some alias; path } ->
      Printf.sprintf "imports %s as %s"
        (string_of_reference path)
        alias
      |> String.escaped
  in
  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
    let error_form = Define ("rego_error", compile_error_condition env') in
    [Comment (Printf.sprintf "package %s"
                (string_of_reference path)
              |> String.escaped)]
    @ List.map imports ~f:(fun x -> Comment (compile_import x))
    @ forms
    @ [error_form]