Source file spilling.ml

(* elpi: embedded lambda prolog interpreter                                  *)
(* license: GNU Lesser General Public License Version 2.1 or later           *)
(* ------------------------------------------------------------------------- *)
open Elpi_util.Util
open Compiler_data
module C = Constants
open ScopedTerm

type spill = { vars_names : binder list; expr : t }
type spills = spill list

let eat ~(type_abbrevs: TypeAssignment.type_abbrevs) args ty =
  let ty = TypeAssignment.deref ty in
  let rec aux args ty =
    match args with
    | [] -> ty
    | _ :: args ->
        match TypeAssignment.look ~type_abbrevs ty with
        | TypeAssignment.Arr (_,Elpi_parser.Ast.Structured.Variadic, _, t) -> if args = [] then aux args t else ty
        | TypeAssignment.Arr (_,Elpi_parser.Ast.Structured.NotVariadic, _, t) -> aux args t
        | _ -> anomaly ("eat args: " ^ TypeAssignment.(show @@ Val ty))
  in
  TypeAssignment.create @@ aux args ty

let is_prop ~type_abbrevs x =
  match TypeAssignment.args_missing_to_prop ~type_abbrevs x with
  | Some n -> List.length n = 0
  | None -> false

let mk_global ~types ~loc f l =
  (* TODO: check only builtins *)
  let s = Symbol.make_builtin f in
  let f_ty = TypingEnv.(resolve_symbol s types).ty |> (fun x -> TypeAssignment.apply x l) |> TypeAssignment.create in
  { scope = Scope.mkResolvedGlobal types s; name = f; ty = f_ty; loc }

let pif_w_name_ty ~types ({ ty; loc } : 'a w_ty_loc) : const =
  mk_global ~types ~loc F.pif [TypeAssignment.deref ty]
let pif_ty ~types (ty : binder) = (pif_w_name_ty ~types ty).ty
let pif_arg_ty ~types ty =
  match TypeAssignment.deref @@ pif_ty ~types ty with
  | TypeAssignment.Arr(_,_,x,_) -> TypeAssignment.create x
  | _ -> assert false

let mk_loc ~loc ~ty it = { ty; it; loc }

(* TODO store the types in Main *)
let add_spilled ~types (l : spill list) t =
  if l = [] then t
  else
    List.fold_right
      (fun { expr; vars_names } t ->
        vars_names
        (* Format.eprintf "fold %a\n" ScopedTerm.pretty t; *)
         |> List.fold_left (fun t (v : binder) ->
            mk_loc ~loc:t.loc ~ty:t.ty
            @@ App (mk_global ~loc:v.loc ~types F.sigmaf [TypeAssignment.deref v.ty],
              [mk_loc ~loc:t.loc ~ty:TypeAssignment.(create (Arr (MRef (MutableOnce.make F.dummyname), NotVariadic,deref v.ty,deref t.ty)))
              @@ Lam(Some v,None,t)
              ])
          )
        (mk_loc ~loc:t.loc ~ty:TypeAssignment.(create (Prop Elpi_parser.Ast.Structured.Function))
        @@ App (mk_global ~types ~loc:t.loc F.andf [], [expr; t ])))
      l t

let mkApp n l = App (n, l)

let is_symbol ~types b = function
| Scope.Global { resolved_to = x } ->
  begin match SymbolResolver.resolved_to types x with
  | Some s  -> TypingEnv.same_symbol types s b
  | _ -> anomaly "unresolved global symbol"
  end
| _ -> false



let app ~type_abbrevs ~types t args =
  if args = [] then t
  else
    let rec aux { loc; it; ty } : t =
      let it, ty =
        match it with
        | App (({ scope = s } as n), xs) when is_symbol ~types Elpi_runtime.Data.Global_symbols.and_ s -> mkApp n (aux_last xs), eat ~type_abbrevs args ty
        | Impl (b, l, s, t) -> Impl (b,l, s, aux t), ty
        | App (n, xs) -> mkApp n (xs @ args), eat ~type_abbrevs args ty
        | UVar(c,l) -> UVar (c,l @ args), eat ~type_abbrevs args ty
        | Discard _ | Lam (_, _, _) | CData _ | Spill (_, _) | Cast (_, _) -> assert false
      in  
        mk_loc ~loc ~ty it
    and aux_last = function [] -> assert false | [ x ] -> [ aux x ] | x :: xs -> x :: aux_last xs in
    aux t

let mk_spilled ~loc ~ty args n : (binder * t) list =
  (* builds the type of the spilled variables, all variables has same type *)
  (* let builf_head_ty tgt_ty =
    let rec aux = function
      | [] -> tgt_ty
      | ScopedTerm.{ ty } :: tl ->
          TypeAssignment.(Arr (MRef (MutableOnce.make F.dummyname), NotVariadic, deref ty, aux tl))
    in
    TypeAssignment.create (aux ctx)
  in *)
  let rec aux n ty =
    if n = 0 then []
    else
    let f =
      incr args;
      F.from_string (Printf.sprintf "%%spill%d" !args)
    in
    let built_tm ty =
      let hd_ty = TypeAssignment.create ty in
      mk_loc ~loc ~ty:hd_ty @@ App (mk_bound_const ~lang:elpi_language f ~loc ~ty:hd_ty,[])
    in
      match ty with
      | TypeAssignment.Arr (_, _, l, r) -> (mk_binder ~lang:elpi_language f ~loc ~ty:(TypeAssignment.create l), built_tm l) :: aux (n-1) r
      | UVar r when MutableOnce.is_set r -> aux n (TypeAssignment.deref r)
      | _ -> anomaly "type abbreviations and spilling, not implemented"
  in
  aux n ty

(* Replaces bound-name references (as generated by mk_spilled) with the given
   terms. *)
let rec apply_subst ~subst = function
  | App ({ scope = Bound l; name = f }, []) when l = elpi_language && F.Map.mem f subst -> (F.Map.find f subst).it
  | App (n, xs) -> App (n, List.map (apply_subst_loc ~subst) xs)
  | UVar (c, xs) -> UVar (c, List.map (apply_subst_loc ~subst) xs)
  | Lam (n, o, t) -> Lam (n, o, apply_subst_loc ~subst t)
  | Impl (d, l, t1, t2) -> Impl (d, l, apply_subst_loc ~subst t1, apply_subst_loc ~subst t2)
  | Cast (t, e) -> Cast (apply_subst_loc ~subst t, e)
  | Spill (t, i) -> Spill (apply_subst_loc ~subst t, i)
  | CData _ | Discard _ as x -> x
and apply_subst_loc ~subst ({ it } as t) = { t with it = apply_subst ~subst it }

(* barendregt_convention (naive implementation) *)
let rec bc ctx t =
  match t with
  | Lam (None, o, t) -> Lam (None, o, bc_loc ctx t)
  | Lam (Some { scope = l; name = c; loc; ty = tya }, o, t) when List.mem (c, l) ctx ->
      let d = fresh_bound () in
      bc ctx (Lam (Some { scope = l; name = d; loc; ty = tya }, o, rename_loc l c d t))
  | Lam ((Some { scope = l; name =  c } as abs), o, t) -> Lam (abs, o, bc_loc ((c, l) :: ctx) t)
  | Impl (b, bl, t1, t2) -> Impl (b, bl, bc_loc ctx t1, bc_loc ctx t2)
  | Cast (t, ty) -> Cast (bc_loc ctx t, ty)
  | Spill (t, i) -> Spill (bc_loc ctx t, i)
  | App (hd, xs) -> App (hd, List.map (bc_loc ctx) xs)
  | Discard _ | UVar _ | CData _ -> t

and bc_loc ctx { loc; ty; it } = { loc; ty; it = bc ctx it }

(* let not_from_pi (_,b) = b = false *)
(* let from_pi (_,b) = b = true *)

let rec apply what v = function
  | App ({ scope = Bound l; name = f; ty = hd_ty } as n, xs) when l = elpi_language && List.exists (fun f' -> F.equal f f'.name) what ->
      App ({ n with ty = TypeAssignment.(create @@ Arr (MRef (MutableOnce.make F.dummyname), NotVariadic, deref v.ty,deref hd_ty))}, v :: xs)
  | App(f,xs) -> App(f, smart_map (apply_loc what v) xs)
  | Lam(n,o,t) -> Lam(n,o,apply_loc what v t)
  | Impl(d,l,t1,t2) -> Impl(d,l,apply_loc what v t1,apply_loc what v t2)
  | Cast(t,e) -> Cast(apply_loc what v t,e)
  | Spill _ -> assert false
  | CData _ | Discard _ | UVar _ as x -> x
and apply_loc what v { loc; ty; it } = { loc; ty; it = apply what v it }

let apply_loc what v t =
  (* Format.eprintf "apply %a to %a in %a\n" (pplist (fun fmt (s,_,_) -> Format.fprintf fmt "%s" s) " ") what ScopedTerm.pretty v ScopedTerm.pretty t; *)
  let t = apply_loc what v t in
  (* Format.eprintf "apply=%a\n" ScopedTerm.pretty t; *)
  t

let rec spill ~type_abbrevs ~types ?(extra = 0) args ({ loc; ty; it } as t) : spills * t list =
  (* Format.eprintf "@[<hov 2>spill %a :@ %a@]\n" pretty t TypeAssignment.pretty (TypeAssignment.deref ty); *)
  match it with
  | CData _ | Discard _ -> ([], [ t ])
  | Cast (t, _) -> spill ~types ~type_abbrevs args t
  | Spill (t, { contents = NoInfo }) -> assert false (* no type checking *)
  | Spill (t, { contents = Phantom _ }) -> assert false (* escapes type checker *)
  | Spill (t, { contents = Main n }) ->
      let ty = t.ty in
      (* Format.eprintf "Spilling of %a with ty %a requires %d slots@." ScopedTerm.pretty_ it TypeAssignment.pretty_mut_once (UVar ty) n; *)
      let vars_names, vars = List.split @@ mk_spilled ~loc ~ty:(TypeAssignment.deref ty) args n in
      let t = app ~type_abbrevs ~types t vars in
      let spills, t = spill1 ~types ~type_abbrevs ~extra:(List.length vars_names) args t in
      (* Format.eprintf "Spilled %a@." ScopedTerm.pretty t; *)
      (spills @ [ { vars_names; expr = t } ], vars)
  (* globals and builtins *)
  | App (({ scope = s } as hd), [{ it = Lam (Some v, o, t); loc = tloc; ty = tty }]) when is_symbol ~types Elpi_runtime.Data.Global_symbols.pi s ->
      let spilled, t = spill1 ~types ~type_abbrevs args t in
      ([], [ { loc; ty; it = App (hd, [{ it = Lam (Some v, o, add_spilled ~types spilled t); loc = tloc; ty = tty }]) } ])
  | App (({ scope = s } as hd), [{ it = Lam (Some v, o, t); loc = tloc; ty = tty }]) when is_symbol ~types Elpi_runtime.Data.Global_symbols.sigma s ->
      (* not to be put in scope of spills *)
      let spilled, t = spill1 ~types ~type_abbrevs args t in
      ([], [ { loc; ty; it = App (hd, [{ it = Lam (Some v, o, add_spilled ~types spilled t); loc = tloc; ty = tty }]) } ])
  | App (({ scope = s; ty = hty } as hd), xs) ->

      let mk_eta_var () = incr args; F.from_string @@ Format.asprintf "%%eta%d" !args in
      let last = if is_symbol ~types Elpi_runtime.Data.Global_symbols.and_ s then List.length xs else -1 in
      let spills, args =
        List.split @@ List.mapi (fun i -> spill ~types ~type_abbrevs ~extra:(if i = last then extra else 0) args) xs
      in
      let args = List.flatten args in
      let spilled = List.flatten spills in
      let it = App (hd, args) in

      let ty = eat ~type_abbrevs args hty in 
      if spilled = [] then
        (spilled, [ { it; loc; ty } ])
      else
        begin match TypeAssignment.args_missing_to_prop ~type_abbrevs ty with
        | None -> (spilled, [ { it; loc; ty } ])
        | Some missing ->
            let rec mk_lam l t =
              match l with
              | [] -> t
              | (v,ty) :: vs -> {loc;ty;it = Lam(Some v,None,mk_lam vs t)} in
            let missing_vars = List.map (fun (ty,arrow) -> 
                let v = mk_eta_var () in
                (mk_binder ~lang:elpi_language v ~loc ~ty, arrow)) missing in
            let missing_args = List.map (fun (v,_) -> { ty; loc; it = App(bind_const v, []) }) missing_vars in
            let t = { it; loc; ty } in
            let t = mk_lam missing_vars @@ add_spilled ~types spilled (app ~type_abbrevs ~types t missing_args) in
            ([], [ t ])
        end
  (* TODO: positive/negative postion, for now we assume :- and => are used in the obvious way *)
  | Impl (R2L, _,_,_) -> 
      let t =
        let premise, head_vars = spill_premise ~type_abbrevs ~types args t in
        List.fold_right (fun (v : binder) t ->
            mk_loc ~loc:premise.loc ~ty:(pif_ty ~types v) @@
              App (pif_w_name_ty ~types v, [ 
                mk_loc ~loc:premise.loc ~ty:(pif_arg_ty ~types v) @@ Lam (Some v, None, t) ])
          ) head_vars premise
      in
      ([],[t])
  | Impl ((L2R|L2RBang) as kind, l, premise, conclusion) ->
      let pi_closed_premise =
        let premise, head_vars = spill_premise ~type_abbrevs ~types args premise in
        List.fold_right (fun (v : binder) t ->
            mk_loc ~loc:premise.loc ~ty:(pif_ty ~types v) @@
              App (pif_w_name_ty ~types v, [ 
                mk_loc ~loc:premise.loc ~ty:(pif_arg_ty ~types v) @@ Lam (Some v, None, t) ])
          ) head_vars premise
      in
      let spilled_conclusion =
        let spilled, conclusion = spill1 ~types ~type_abbrevs ~extra args conclusion in
        if spilled <> [] then anomaly ~loc "conclusion of ==> leaks spilled calls";
        conclusion in
      ([], [ { it = Impl (kind, l, pi_closed_premise, spilled_conclusion); loc; ty } ])
  (* lambda terms *)
  | Lam (None, o, t) ->
      let spills, t = spill1 ~types ~type_abbrevs args t in
      (spills, [ { it = Lam (None, o, t); loc; ty } ])
  | Lam ((Some c as abs), o, t) ->
      let spills, t = spill1 ~types ~type_abbrevs args t in
      let t, spills =
        let { scope = s; name = f; ty } = c in
        map_acc
          (fun t { vars_names; expr } ->
            let bc = mk_loc ~loc ~ty (App(bind_const c,[])) in
            ( apply_loc vars_names bc t,
              {
                vars_names = List.map (fun (v : binder) ->
                   { v with ty = TypeAssignment.(create @@ Arr (MRef (MutableOnce.make F.dummyname), NotVariadic, deref ty,deref v.ty)) }) vars_names;
                expr =
                  mk_loc ~loc ~ty:(pif_ty ~types  c) @@
                    App (pif_w_name_ty ~types c, [mk_loc ~loc ~ty:(pif_arg_ty ~types c) @@ Lam (abs, o, apply_loc vars_names bc expr)]);
              } ))
          t spills
      in
      (spills, [ { it = Lam (abs, o, t); loc; ty } ])
  (* holes *)
  | UVar ({ ty = cty } as c, xs) ->
      let spills, args = List.split @@ List.map (spill ~types ~type_abbrevs args) xs in
      let args = List.flatten args in
      let spilled = List.flatten spills in
      let it = UVar (c, args) in
      let ty = eat ~type_abbrevs args cty in
      if is_prop ~type_abbrevs ty then ([], [ add_spilled ~types spilled { it; loc; ty } ]) else (spilled, [ { it; loc; ty } ])

and spill1 ~type_abbrevs ~types ?extra args ({ loc } as t) =
  let spills, t = spill ~types ~type_abbrevs ?extra args t in
  let t = if List.length t <> 1 then error ~loc "bad spilling" else List.hd t in
  (spills, t)

(* Spills the arguments of top-level application, without post-processing. *)
and spill_top_app_args ~type_abbrevs ~types args t =
  match t.it with
  | App (hd, xs) ->
      let is_input = TypeAssignment.modes_of_ty ~type_abbrevs hd.ty in
      let spills, xs =
        List.split @@ List.mapi (fun i x ->
            let spills, x' = spill ~types ~type_abbrevs args x in
            (if spills <> [] && (match List.nth_opt is_input i with Some b -> b | None -> false) then
               error ~loc:x.loc
                 "Spilling forbidden in input position of a rule's head");
            spills, x')
          xs
      in
      List.flatten spills, { t with it = App (hd, List.flatten xs) }
  | UVar (hd, xs) ->
      let spills, xs = List.split @@ List.map (spill ~types ~type_abbrevs args) xs in
      List.flatten spills, { t with it = UVar (hd, List.flatten xs) }
  | _ -> spill1 ~types ~type_abbrevs args t

(* Spills the hypothesis (premise) of an implication `H => G`. Returns
   the (possibly unchanged) premise together with the placeholders to
   pi-bind, empty when nothing was spilled. *)
and spill_premise ~type_abbrevs ~types args premise =
  let head, body = match premise.it with
    | Impl (R2L, _, h, b) -> h, Some b
    | _ -> premise, None
  in
  let spills_head, head = spill_top_app_args ~type_abbrevs ~types args head in
  let body = Option.map (fun b ->
      let spilled, b = spill1 ~types ~type_abbrevs args b in
      add_spilled ~types spilled b) body
  in
  match spills_head with
  | [] ->
      (match body with
       | None -> premise, []
       | Some body -> { premise with it = Impl (R2L, premise.loc, head, body) }, [])
  | _ ->
      let novars = List.map (fun { expr; _ } -> { vars_names = []; expr }) spills_head in
      let body =
        match body with
        | Some body -> add_spilled ~types novars body
        | None ->
            (match List.rev novars with
             | { expr = last; _ } :: rest_rev -> add_spilled ~types (List.rev rest_rev) last
             | [] -> assert false)
      in
      { premise with it = Impl (R2L, premise.loc, head, body) },
      List.concat_map (fun { vars_names; _ } -> vars_names) spills_head

let spill_toplevel_clause_head ~type_abbrevs ~types args head premise =
  let spills_head, head = spill_top_app_args ~type_abbrevs ~types args head in
  let spilled_to_uvar =
    spills_head |> List.fold_left (fun m { vars_names; _ } ->
      vars_names |> List.fold_left (fun m (v : binder) ->
        incr args;
        let name = F.from_string (Printf.sprintf "%%headspill%d" !args) in
        let uv = mk_loc ~loc:v.loc ~ty:v.ty @@ UVar (mk_uvar name ~loc:v.loc ~ty:v.ty, []) in
        F.Map.add v.name uv m
       ) m
      ) F.Map.empty in
  let head = apply_subst_loc ~subst:spilled_to_uvar head in
  let spills_head =
    List.map (fun { expr; _ } ->
      { vars_names = []; expr = apply_subst_loc ~subst:spilled_to_uvar expr }) spills_head in
  match premise, spills_head with
  | _, [] -> head, premise
  | Some premise, _ -> head, Some (add_spilled ~types spills_head premise)
  | None, _ ->
      match List.rev spills_head with
      | { expr = last; _ } :: rest_rev -> head, Some (add_spilled ~types (List.rev rest_rev) last)
      | [] -> assert false

let fresh = ref 0

let rec remove_top_sigmas ~types t =
  match t.it with
  | App ({ scope = s }, [_]) when is_symbol ~types Elpi_runtime.Data.Global_symbols.pi s -> t
  | App ({ scope = s } as n,xs) when is_symbol ~types Elpi_runtime.Data.Global_symbols.and_ s ->
      { t with it = App(n, smart_map (remove_top_sigmas ~types) xs) }
  | Impl(x,l,t1,t2) -> { t with it = Impl(x,l,t1,remove_top_sigmas ~types t2) }
  | App ({ scope = s }, [{ it = Lam(Some { name = vn; ty = vty; loc=vloc },_,{ loc;ty }); } as b]) when is_symbol ~types Elpi_runtime.Data.Global_symbols.sigma s ->
      let vn_fresh = F.from_string @@ Printf.sprintf "%s_%d" (F.show vn) !fresh in
      incr fresh;
      remove_top_sigmas ~types { loc; ty; it = ScopedTerm.beta b [{ ty = vty; loc; it = UVar(mk_uvar vn_fresh ~loc:vloc ~ty:vty,[]) }] }
  | _ -> t

let main_goal ~type_abbrevs ~types t =
  let args = ref 0 in
  let t = bc_loc [] t in
  let spills, ts = spill ~type_abbrevs ~types args (bc_loc [] t) in
  if spills <> [] then error ~loc:t.loc "Cannot place spilled expression";
  let ts = List.map (remove_top_sigmas ~types) ts in
  match ts with
  | [ t ] -> t
  |  _ -> assert false

let main_clause ~type_abbrevs ~types t =
  let args = ref 0 in
  let t = bc_loc [] t in
  match t.it with
  | Impl (R2L, l, head, premise) ->
      let head, premise =
        match spill_toplevel_clause_head ~type_abbrevs ~types args head (Some premise) with
        | head, Some premise -> head, premise
        | _, None -> assert false (* premise was Some, so it stays Some *) in
      let spilled, premise = spill1 ~types ~type_abbrevs args premise in
      if spilled <> [] then error ~loc:t.loc "Cannot place spilled expression";
      remove_top_sigmas ~types { t with it = Impl (R2L, l, head, premise) }
  | _ ->
      match spill_toplevel_clause_head ~type_abbrevs ~types args t None with
      | head, None -> head
      | head, Some premise -> 
          remove_top_sigmas ~types { t with it = Impl (R2L, t.loc, head, premise) }