Source file ppx_deriving_jsoo_lib.ml

open Ppxlib
open Ast_builder.Default
module Common = Common
module Jsoo_type = Jsoo_type
module Jsoo_conv = Jsoo_conv
module Ppx_js = Ppx_js
open Common

module Base = struct
  let jsoo_types ?modules ?enum ?camel ?snake ~case ~rm_prefix ~mut l =
    let l = List.map (fun t ->
        let loc = t.ptype_loc in
        let name = t.ptype_name.txt in
        let name_js = mkl ~loc @@ jsoo_name ?modules name in
        let params = t.ptype_params in
        let ct_or_t, ct_ext =
          Jsoo_type.declaration_of_type_kind ?rm_prefix ?modules ?enum ?camel ?snake ~case ~mut t in
        match ct_or_t with
        | CT expr ->
          ct_ext @ [class_infos ~loc ~virt:Concrete ~params ~name:name_js ~expr], None, true
        | TT c -> ct_ext, Some (
            type_declaration ~loc ~name:name_js ~params ~cstrs:[]
              ~kind:Ptype_abstract ~manifest:(Some c) ~private_:Public), false
      ) l in
    List.flatten @@ List.map (fun (x, _, _) -> x) l,
    List.filter_map (fun (_, x, _) -> x) l,
    List.map (fun (_, _, x) -> x) l

  let jsoo_convs ?modules ?enum ~case ~rm_prefix ~mut ?camel ?snake ?remove_undefined l =
    List.map (fun t ->
        let cv = Jsoo_conv.conv_expressions ?rm_prefix ?modules ?enum ?camel ?snake ?remove_undefined ~case ~mut  t in
        t, cv) l

  let jsoo_str ~loc ~rec_flag ~case ~rm_prefix ~recursive ~mut ?modules ?enum ?camel ?snake ?remove_undefined l =
    let cts, ts, are_class_type = jsoo_types ~case ~rm_prefix ~mut ?modules ?enum ?camel ?snake l in
    let convs = jsoo_convs ~case ~rm_prefix ~mut ?modules ?enum ?camel ?snake ?remove_undefined l in
    let vals, conv = List.split @@ List.map2 (fun (t, cv) is_class_type ->
        let loc = t.ptype_loc in
        let name = t.ptype_name.txt in
        let _name_js, name_to, name_of, name_conv =
          jsoo_name ?modules name, jsoo_name_to ?modules name, jsoo_name_of ?modules name, jsoo_name_conv ?modules name in
        let to_sig, of_sig, conv_sig = Jsoo_conv.conv_signatures ~is_class_type t in
        let acc = List.flatten @@ List.map (fun ({e_to; e_of; _}, (name_to, name_of)) ->
            [ value_binding ~loc ~pat:(pvar ~loc name_to) ~expr:e_to;
              value_binding ~loc ~pat:(pvar ~loc name_of) ~expr:e_of ]) cv.e_acc in
        acc @ [ value_binding ~loc ~pat:(ppat_constraint ~loc (pvar ~loc name_to) to_sig)
                  ~expr:cv.e_to;
                value_binding ~loc ~pat:(ppat_constraint ~loc (pvar ~loc name_of) of_sig)
                  ~expr:cv.e_of ],
        value_binding ~loc ~pat:(ppat_constraint ~loc (pvar ~loc name_conv) conv_sig)
          ~expr:(pexp_tuple ~loc [evar ~loc name_to; evar ~loc name_of])
      ) convs are_class_type in
    let rec_flag_type = match recursive with
      | None -> rec_flag | Some true -> Recursive | Some false -> Nonrecursive in
    let vals = List.flatten vals in
    let rec_flag_conv = match recursive, List.length vals with
      | Some true, _ -> Recursive
      | Some false, _ -> Nonrecursive
      | None, i when i <= 2 -> Nonrecursive
      | _ -> Recursive in
    pstr_class_type ~loc cts :: (match ts with [] -> [] | _ -> [ pstr_type ~loc rec_flag_type ts ])
    @ [ pstr_value ~loc rec_flag_conv vals; pstr_value ~loc Nonrecursive conv ]

  let jsoo_sig ~loc ~rec_flag ~case ~rm_prefix ~mut ?modules ?enum ?camel ?snake l =
    let cts, ts, are_class_type = jsoo_types ~case ~rm_prefix ~mut ?modules ?enum ?camel ?snake l in
    let sigs = List.flatten @@ List.map2 (fun t is_class_type ->
        let loc = t.ptype_loc in
        let name = t.ptype_name.txt in
        let _name_js, name_to, name_of, name_conv =
          jsoo_name ?modules name, jsoo_name_to ?modules name, jsoo_name_of ?modules name, jsoo_name_conv ?modules name in
        let to_sig, of_sig, conv_sig = Jsoo_conv.conv_signatures ~is_class_type t in
        [ value_description ~loc ~name:{txt=name_to; loc} ~type_:to_sig ~prim:[];
          value_description ~loc ~name:{txt=name_of; loc} ~type_:of_sig ~prim:[];
          value_description ~loc ~name:{txt=name_conv; loc} ~type_:conv_sig ~prim:[];
        ]) l are_class_type in
    let s = psig_class_type ~loc cts :: (
          match ts with [] -> [] | _ -> [ psig_type ~loc rec_flag ts ]) in
    s @ List.map (psig_value ~loc) sigs

  let jsoo_types_str ~loc ~rec_flag ~case ~rm_prefix ~recursive ~mut ?modules ?enum ?camel ?snake l =
    let rec_flag = match recursive with
      | None -> rec_flag | Some true -> Recursive | Some false -> Nonrecursive in
    let cts, ts, _ = jsoo_types ~case ~rm_prefix ~mut ?modules ?enum ?camel ?snake l in
    pstr_class_type ~loc cts :: (
      match ts with [] -> [] | _ -> [ pstr_type ~loc rec_flag ts ])

  let jsoo_convs_str ~loc ~case ~rm_prefix ~recursive ~mut ?modules ?enum ?camel ?snake ?remove_undefined l =
    let l = jsoo_convs ~case ~rm_prefix ~mut ?modules ?enum ?camel ?snake ?remove_undefined l in
    let vals, conv = List.split @@ List.map (fun (t, cv) ->
        let loc = t.ptype_loc in
        let name = t.ptype_name.txt in
        let name_to, name_of, name_conv =
          jsoo_name_to ?modules name, jsoo_name_of ?modules name, jsoo_name_conv ?modules name in
        let acc = List.flatten @@ List.map (fun ({e_to; e_of; _}, (name_to, name_of)) ->
            [ value_binding ~loc ~pat:(pvar ~loc name_to) ~expr:e_to;
              value_binding ~loc ~pat:(pvar ~loc name_of) ~expr:e_of ]) cv.e_acc in
        acc @ [ value_binding ~loc ~pat:(pvar ~loc name_to) ~expr:cv.e_to;
                value_binding ~loc ~pat:(pvar ~loc name_of) ~expr:cv.e_of ],
        value_binding ~loc ~pat:(pvar ~loc name_conv)
          ~expr:(pexp_tuple ~loc [evar ~loc name_to; evar ~loc name_of])) l in
    let vals = List.flatten vals in
    let rec_flag = match recursive, List.length vals with
      | Some true, _ -> Recursive
      | Some false, _ -> Nonrecursive
      | None, i when i <= 2 -> Nonrecursive
      | _ -> Recursive in
    [ pstr_value ~loc rec_flag vals; pstr_value ~loc Nonrecursive conv ]
end

module Ppx_deriver = struct
  open Base

  let str_jsoo_types ~loc ~path:_ (rec_flag, types) case force_debug rm_prefix
      recursive mut modules enum camel snake _remove_undefined =
    if fake then []
    else
      let str = jsoo_types_str ~loc ~rec_flag ~case ~rm_prefix ~recursive ~mut
          ?modules ?enum ~camel ~snake types in
      debug ~force:force_debug "%s" (str_of_structure str);
      str

  let sig_jsoo_types ~loc ~path:_ (rec_flag, l) case rm_prefix recursive mut
      modules enum camel snake =
    if fake then []
    else
      let rec_flag = match recursive with
        | None -> rec_flag | Some true -> Recursive | Some false -> Nonrecursive in
      let cts, ts, _ = jsoo_types ~case ~rm_prefix ~mut ?modules ?enum ~camel ~snake l in
      let s = psig_class_type ~loc cts :: (
          match ts with [] -> [] | _ -> [ psig_type ~loc rec_flag ts ]) in
      s

  let str_jsoo_convs ~loc ~path:_ (_rec_flag, l) case force_debug rm_prefix
      recursive mut modules enum camel snake remove_undefined =
    if fake then []
    else
      let str = jsoo_convs_str ~loc ~case ~rm_prefix ~recursive ~mut ?modules
          ?enum ~camel ~snake ~remove_undefined l in
      debug ~force:force_debug "%s" (str_of_structure str);
      Ppx_js.transform#structure str

  let str_gen ~loc ~path:_ (rec_flag, l) case force_debug rm_prefix recursive mut
      modules enum camel snake remove_undefined =
    if fake then []
    else
      let str = jsoo_str ~loc ~rec_flag ~case ~rm_prefix ~recursive ~mut ?modules
          ?enum ~camel ~snake ~remove_undefined l in
      debug ~force:force_debug "%s" (str_of_structure str);
      Ppx_js.transform#structure str

  let sig_gen ~loc ~path:_ (rec_flag, l) case rm_prefix _recursive mut modules enum camel snake =
    if fake then []
    else
      let s = jsoo_sig ~loc ~rec_flag ~case ~rm_prefix ~mut ?modules ?enum ~camel ~snake l in
      Pprintast.signature Format.err_formatter s;
      Ppx_js.transform#signature s

  let ebool t =
    let f = Deriving.Args.to_func t in
    Deriving.Args.of_func (fun ctx loc x k ->
        match x.pexp_desc with
        | Pexp_construct ({txt=Lident "true"; _}, _) ->
          f ctx loc true k
        | Pexp_construct ({txt=Lident "false"; _}, _) ->
          f ctx loc false k
        | _ -> Location.raise_errorf ~loc "wrong boolean argument")

  let construct_pair t =
    let f = Deriving.Args.to_func t in
    Deriving.Args.of_func (fun ctx loc x k ->
        match x.pexp_desc with
        | Pexp_tuple [ {pexp_desc=Pexp_construct ({txt=id1; _}, None); _};
                       {pexp_desc=Pexp_construct ({txt=id2; _}, None); _} ] ->
          f ctx loc (Longident.name id1, Longident.name id2) k
        | _ -> Location.raise_errorf ~loc "wrong ident pair argument")

  let enum t =
    let f = Deriving.Args.to_func t in
    Deriving.Args.of_func (fun ctx loc x k ->
        match x.pexp_desc with
        | Pexp_construct ({txt=Lident "String"; _}, _) ->
          f ctx loc `string k
        | Pexp_construct ({txt=Lident "Int"; _}, _) ->
          f ctx loc `int k
        | _ ->
          f ctx loc `string k)

  let args_str = Deriving.Args.(
      empty
      +> flag "case"
      +> flag "debug"
      +> arg "remove_prefix" (ebool __)
      +> arg "recursive" (ebool __)
      +> flag "mut"
      +> arg "modules" (elist (construct_pair __))
      +> arg "enum" (enum __)
      +> flag "camel"
      +> flag "snake"
      +> flag "remove_undefined"
    )

  let args_sig = Deriving.Args.(
      empty
      +> flag "case"
      +> arg "remove_prefix" (ebool __)
      +> arg "recursive" (ebool __)
      +> flag "mut"
      +> arg "modules" (elist (construct_pair __))
      +> arg "enum" (enum __)
      +> flag "camel"
      +> flag "snake"
    )

  let main () =
    let str_type_decl = Deriving.Generator.make args_str str_gen in
    let sig_type_decl = Deriving.Generator.make args_sig sig_gen in
    Deriving.ignore @@ Deriving.add jsoo ~str_type_decl ~sig_type_decl;
    let str_type_decl = Deriving.Generator.make args_str str_jsoo_types in
    let sig_type_decl = Deriving.Generator.make args_sig sig_jsoo_types in
    Deriving.ignore @@ Deriving.add "jsoo_type" ~str_type_decl ~sig_type_decl;
    let str_type_decl = Deriving.Generator.make args_str str_jsoo_convs in
    Deriving.ignore @@ Deriving.add "jsoo_conv" ~str_type_decl
end

module Deriver_exe = struct
  open Base

  type options = {
    case : bool;
    remove_prefix : bool option;
    recursive : bool option;
    mut : bool;
    debug : bool;
    modules : (string * string) list option;
    enum : [ `string | `int ] option;
    remove_undefined : bool;
  }

  let options = {
    case=false; remove_prefix=None; recursive=None;
    mut=false; debug=false; modules=None; enum=None; remove_undefined=false}

  let get_bool s e = match e.pexp_desc with
    | Pexp_construct ({txt=Lident "true"; _}, _) -> true
    | Pexp_construct ({txt=Lident "false"; _}, _) -> false
    | _ -> Location.raise_errorf ~loc:e.pexp_loc "wrong boolean argument for %S" s

  let rec get_list acc e = match e.pexp_desc with
    | Pexp_construct ({ txt = Lident "[]"; _ }, None) ->
      List.rev acc
    | Pexp_construct ({ txt = Lident "::"; _ }, Some arg) -> begin
        match arg.pexp_desc with
        | Pexp_tuple [hd; tl] -> get_list (hd :: acc) tl
        | _ ->  Location.raise_errorf ~loc:arg.pexp_loc "wrong list argument"
      end
    | _ -> Location.raise_errorf ~loc:e.pexp_loc "wrong list argument"

  let get_construct_pair s e =
    match e.pexp_desc with
    | Pexp_tuple [ {pexp_desc=Pexp_construct ({txt=id1; _}, None); _}; {pexp_desc=Pexp_construct ({txt=id2; _}, None); _} ] ->
      (Longident.name id1, Longident.name id2)
    | _ ->
      Location.raise_errorf ~loc:e.pexp_loc "wrong ident pair for %S" s

  let get_enum e = match e.pexp_desc with
    | Pexp_construct ({txt=Lident "String"; _}, _) -> Some `string
    | Pexp_construct ({txt=Lident "Int"; _}, _) -> Some `int
    | _ -> Some `string

  let get_options = function
    | {pexp_desc=Pexp_record (l, _); _} ->
      let l = List.filter_map (function ({txt=Lident s; _}, e) -> Some (s, e) | _ -> None) l in
      let opt = List.fold_left (fun acc (s, e) -> match s with
          | "case" -> {acc with case = true}
          | "remove_prefix" -> {acc with remove_prefix = Some (get_bool s e)}
          | "recursive" -> {acc with recursive = Some (get_bool s e)}
          | "mut" -> {acc with mut = true}
          | "debug" -> {acc with debug = true}
          | "modules" -> {acc with modules = Some (List.map (get_construct_pair s) (get_list [] e))}
          | "enum" -> {acc with enum = get_enum e}
          | "remove_undefined" -> {acc with remove_undefined = true}
          | _ -> Location.raise_errorf ~loc:e.pexp_loc "argument %S not handled" s
        ) options l in
      Some opt
    | _ ->
      Some options

  let filter l =
    let t = List.hd @@ List.rev l in
    let rec iter = function
      | [] -> None
      | a :: tl ->
        if a.attr_name.txt <> "deriving" then iter tl
        else match a.attr_payload with
          | PStr str ->
            let rec aux = function
              | [] -> iter tl
              | it :: str ->
                match it.pstr_desc with
                | Pstr_eval ({pexp_desc=Pexp_ident {txt=Lident "jsoo"; _}; _}, _) ->
                  Some options
                | Pstr_eval ({pexp_desc=Pexp_apply ({pexp_desc=Pexp_ident {txt=Lident "jsoo"; _}; _}, [Nolabel, e]); _}, _) ->
                  get_options e
                | Pstr_eval ({pexp_desc=Pexp_tuple l; _}, _) ->
                  let rec aux2 = function
                    | [] -> aux str
                    | h :: etl -> match h with
                      | {pexp_desc=Pexp_ident {txt=Lident "jsoo"; _}; _} ->
                        Some options
                      | {pexp_desc=Pexp_apply ({pexp_desc=Pexp_ident {txt=Lident "jsoo"; _}; _}, [Nolabel, e]); _} ->
                        get_options e
                      | _ -> aux2 etl in
                  aux2 l
                | _ -> aux str in
            aux str
          | _ -> iter tl in
    iter t.ptype_attributes

  let process ~loc ?(kind=`all) ~rec_flag ~case ~rm_prefix ~recursive ~mut ?modules ?enum ?remove_undefined l =
    match kind with
    | `all -> jsoo_str ~loc ~rec_flag ~case ~rm_prefix ~recursive ~mut ?modules ?enum ?remove_undefined l
    | `typ -> jsoo_types_str ~loc ~rec_flag ~case ~rm_prefix ~recursive ~mut ?modules ?enum l
    | `conv -> jsoo_convs_str ~loc ~case ~rm_prefix ~recursive ~mut ?modules ?enum ?remove_undefined l

  let rec derive_str ?(loc=Location.none) ?kind name str =
    let str = List.fold_left (fun acc it ->
        match it.pstr_desc with
        | Pstr_type (rec_flag, l) ->
          begin match filter l with
            | None -> acc
            | Some o ->
              let str = process ~loc:it.pstr_loc ?kind ~rec_flag ~case:o.case
                  ~rm_prefix:o.remove_prefix ~recursive:o.recursive ~mut:o.mut
                  ?modules:o.modules ?enum:o.enum ~remove_undefined:o.remove_undefined l in
              Common.debug ~force:o.debug "%s" (Common.str_of_structure str);
              acc @ (Ppx_js.transform#structure str)
          end
        | Pstr_module p ->
          begin match p.pmb_name.txt with
            | None -> acc
            | Some name ->
              begin match p.pmb_expr.pmod_desc with
                | Pmod_structure str ->
                  let str = derive_str ?kind name str in
                  if List.length str <= 1 then acc
                  else
                    let loc = p.pmb_expr.pmod_loc in
                    acc @ [
                      pstr_module ~loc @@
                      module_binding ~loc ~name:{txt=Some name; loc}
                        ~expr:(pmod_structure ~loc str) ]
                | _ -> acc
              end
          end
        | Pstr_attribute a ->
          if a.attr_name.txt <> "jsoo" then acc
          else begin match a.attr_payload with
            | PStr str -> acc @ Ppx_js.transform#structure str
            | _ -> acc end
        | _ -> acc) [] str in
    pstr_open ~loc (open_infos ~loc ~override:Fresh ~expr:(pmod_ident ~loc (Common.llid ~loc name))) :: str

  let main () =
    let filename = ref None in
    let outfile = ref None in
    let kind = ref `all in
    let specs = [
      "--output", Arg.String (fun s -> outfile := Some s), "output file";
      "-o", Arg.String (fun s -> outfile := Some s), "alias of --output";
      "--type-only", Arg.Unit (fun () -> kind := `typ), "only output type/class type";
      "--conv-only", Arg.Unit (fun () -> kind := `conv), "only output conversion function";
    ] in
    let usage = "derive_jsoo <options> <filename>" in
    Arg.parse specs (fun f -> filename := Some f) usage;
    match !filename with
    | None -> Arg.usage specs usage
    | Some f ->
      let name = String.capitalize_ascii @@ Filename.(remove_extension @@ basename f) in
      let ic = open_in f in
      let s = really_input_string ic (in_channel_length ic) in
      let lexbuf = Lexing.from_string s in
      let str = Parse.implementation lexbuf in
      let str = derive_str ~kind:!kind name str in
      match !outfile with
      | None -> Pprintast.structure Format.std_formatter str
      | Some f ->
        let oc = open_out f in
        Pprintast.structure (Format.formatter_of_out_channel oc) str

end