ppx_eliom_client.ml

open Ppxlib
open Ast_helper
open Ppx_eliom_utils

let attribute_of_warning loc s =
  Attr.mk
    {loc; txt = "ocaml.ppwarning"}
    (PStr [Str.eval ~loc (Exp.constant (Const.string s))])

module Pass = struct
  (** {2 Auxiliaries} *)

  (* Replace every escaped identifier [v] with
     [Eliom_client_core.Syntax_helpers.get_escaped_value v] *)
  let map_get_escaped_values expr =
    (object
       inherit Ppxlib.Ast_traverse.map as super

       method! expression e =
         match e.pexp_desc with
         | Pexp_ident {txt; _}
           when Mli.is_escaped_ident @@ Longident.last_exn txt ->
             let loc = e.pexp_loc in
             [%expr Eliom_client_core.Syntax_helpers.get_escaped_value [%e e]]
         | _ -> super#expression e
    end)
      #expression
      expr

  let push_escaped_binding, flush_escaped_bindings =
    let server_arg_ids = ref [] in
    let push gen_id (expr : expression) get_type =
      match
        List.find_opt
          (fun (gen_id', _, _) -> gen_id.txt = gen_id'.txt)
          !server_arg_ids
      with
      | Some (_, _, typ) -> typ
      | None ->
          let typ = get_type () in
          server_arg_ids := (gen_id, expr, typ) :: !server_arg_ids;
          typ
    in
    let flush () =
      let res = List.rev !server_arg_ids in
      server_arg_ids := [];
      res
    in
    push, flush

  let mark_injection, flush_injection =
    let has_injection = ref false in
    let mark () = has_injection := true in
    let flush () =
      let x = !has_injection in
      has_injection := false;
      x
    in
    mark, flush

  let push_client_value_data, flush_client_value_datas =
    let client_value_datas = ref [] in
    let push loc gen_num gen_id expr (args : string Location.loc list) =
      client_value_datas :=
        (loc, gen_num, gen_id, expr, args) :: !client_value_datas
    in
    let flush () =
      let res = List.rev !client_value_datas in
      client_value_datas := [];
      res
    in
    push, flush

  let find_escaped_ident loc id =
    if Mli.exists ()
    then Mli.find_escaped_ident id
    else if Cmo.exists ()
    then Cmo.find_escaped_ident loc
    else [%type: _]

  let find_injected_ident loc id =
    if Mli.exists ()
    then Mli.find_injected_ident id
    else if Cmo.exists ()
    then Cmo.find_injected_ident loc
    else [%type: _]

  let find_fragment loc id =
    if Mli.exists ()
    then Mli.find_fragment id
    else
      match if Cmo.exists () then Cmo.find_fragment loc else None with
      | Some typ -> typ
      | None -> [%type: _]

  let register_client_closures client_value_datas =
    let registrations =
      List.map
        (fun (loc, num, id, expr, args) ->
           let typ = find_fragment loc id in
           let args = List.map Pat.var args in
           let loc = expr.pexp_loc in
           [%expr
             Eliom_client_core.Syntax_helpers.register_client_closure
               [%e str num] (fun [%p pat_args args] : [%t typ] ->
               [%e map_get_escaped_values expr])])
        client_value_datas
    in
    match registrations with
    | [] -> []
    | _ -> [Str.eval (sequence registrations)]

  (* We hoist the body of client fragments to enforce the correct scoping:
     Identifiers declared earlier in the client section should not be
     visible inside the client fragment (unless via escaped value). *)
  let define_client_functions ~loc client_value_datas =
    match client_value_datas with
    | [] -> []
    | _ ->
        let bindings =
          List.map
            (fun (loc, _num, id, expr, args) ->
               let patt = Pat.var id in
               let typ = find_fragment loc id in
               let args = List.map Pat.var args in
               let expr =
                 [%expr fun [%p pat_args args] : [%t typ] -> [%e expr]]
               in
               Vb.mk ~loc patt expr)
            client_value_datas
        in
        [Str.value ~loc Nonrecursive bindings]

  (* For injections *)

  let close_server_section loc =
    [%stri
      let () =
        Eliom_client_core.Syntax_helpers.close_server_section
          [%e eid @@ id_file_hash loc]]

  let may_close_server_section ~no_fragment item =
    if no_fragment then [] else [close_server_section item.pstr_loc]

  let open_client_section loc =
    [%stri
      let () =
        Eliom_client_core.Syntax_helpers.open_client_section
          [%e eid @@ id_file_hash loc]]

  let may_open_client_section loc =
    if flush_injection () then [open_client_section loc] else []

  (** Syntax extension *)

  let client_str item =
    let loc = item.pstr_loc in
    may_open_client_section loc @ [item]

  let server_str no_fragment item =
    register_client_closures (flush_client_value_datas ())
    @ may_close_server_section ~no_fragment item

  let shared_str no_fragment item =
    let loc = item.pstr_loc in
    let client_expr_data = flush_client_value_datas () in
    let op = may_open_client_section loc in
    op
    @ register_client_closures client_expr_data
    @ define_client_functions ~loc client_expr_data
    @ [item]
    @ may_close_server_section ~no_fragment:(no_fragment || op <> []) item

  let fragment ~loc ?typ ~context ~num ~id ~unsafe expr =
    let frag_eid = eid id in
    let escaped_bindings = flush_escaped_bindings () in
    (match typ with
    | Some _ -> ()
    | None when not (Mli.exists () || Cmo.exists ()) -> ()
    | None -> (
      match find_fragment loc id with
      | {ptyp_desc = Ptyp_var _; _} when not unsafe ->
          Location.raise_errorf ~loc
            "The types of client values must be monomorphic from its usage or from its type annotation"
      | _ -> ()));
    push_client_value_data loc num id expr
      (List.map (fun (gen_id, _, _) -> gen_id) escaped_bindings);
    match context, escaped_bindings with
    | `Server, _ ->
        (* We are in a server fragment, this code should always be discarded. *)
        Exp.extension @@ Location.Error.to_extension
        @@ Location.Error.make ~loc ~sub:[] "Eliom: ICE"
    | `Shared, [] -> [%expr [%e frag_eid] ()]
    | `Shared, _ ->
        let bindings =
          List.map
            (fun (gen_id, expr, _) ->
               Vb.mk ~loc:expr.pexp_loc (Pat.var gen_id) expr)
            escaped_bindings
        in
        let args =
          format_args @@ List.map (fun (id, _, _) -> eid id) escaped_bindings
        in
        Exp.let_ ~loc Nonrecursive bindings [%expr [%e frag_eid] [%e args]]

  let check_no_variable =
    (object
       inherit Ppxlib.Ast_traverse.map as super

       method! core_type typ =
         match typ with
         | {ptyp_desc = Ptyp_var _; ptyp_loc = loc; _} ->
             let attr =
               attribute_of_warning loc
                 "The type of this injected value contains a type variable that could be wrongly inferred."
             in
             { typ with
               ptyp_attributes = attr :: typ.ptyp_attributes
             ; ptyp_loc = loc }
         | _ -> super#core_type typ
    end)
      #core_type

  let escape_inject ~loc:loc0 ?ident ~(context : Context.escape_inject) ~id
      ~unsafe expr
    =
    let loc = expr.pexp_loc in
    let frag_eid = eid id in
    let assert_no_variables t = if unsafe then t else check_no_variable t in
    match context with
    (* [%%server [%client ~%( ... ) ] ] *)
    | `Escaped_value _section ->
        let typ =
          push_escaped_binding id expr (fun () ->
            let typ = find_escaped_ident loc0 id in
            let typ = assert_no_variables typ in
            typ)
        in
        [%expr ([%e frag_eid] : [%t typ])]
    (* [%%server ... %x ... ] *)
    | `Injection _section ->
        mark_injection ();
        let typ = find_injected_ident loc0 id in
        let typ = assert_no_variables typ in
        let ident =
          match ident with
          | None -> [%expr None]
          | Some i -> [%expr Some [%e str i]]
        in
        let u, d = Mli.get_injected_ident_info id.txt in
        let es = str ~loc:id.loc (Printf.sprintf "%s%d" u d) in
        [%expr
          (Eliom_client_core.Syntax_helpers.get_injection ?ident:[%e ident]
             ~pos:[%e position loc] [%e es]
           : [%t typ])]

  let shared_sig item = [item]
  let server_sig _ = []
  let client_sig item = [item]
  let prelude _ = []
  let postlude _ = []
end

include Make (Pass)

let () =
  Ppxlib.Driver.register_transformation ~impl:mapper#structure
    ~intf:mapper#signature "ppx_eliom_client"