Source file utils.ml

open Ppxlib
open Ast_builder.Default

let verbose = match Sys.getenv_opt "PPX_ENCODING_DEBUG" with
  | None | Some "0" | Some "false" | Some "no" -> 0
  | Some s ->
    match s with
    | "true" -> 1
    | s -> match int_of_string_opt s with
      | Some i -> i
      | None -> 0

let fake = match Sys.getenv_opt "PPX_ENCODING_FAKE" with
  | Some "1" | Some "true" | Some "yes" -> true
  | _ -> false

let remove_prefix_var = match Sys.getenv_opt "PPX_ENCODING_RM_PREFIX" with
  | Some "false" | Some "0" -> false
  | _ -> true

let debug ?(v=1) ?(force=false) fmt =
  if force || verbose >= v then Format.ksprintf (fun s -> Format.eprintf "%s@." s) fmt
  else Printf.ifprintf () fmt

let json_encoding_dot = ref "Json_encoding."
let json_encoding_tmp = ref !json_encoding_dot

let wrap = function
  | None -> json_encoding_tmp := !json_encoding_dot
  | Some "" -> json_encoding_tmp := !json_encoding_dot; json_encoding_dot := ""
  | Some s -> json_encoding_tmp := !json_encoding_dot; json_encoding_dot := String.capitalize_ascii s ^ "."

let unwrap () =
  json_encoding_dot := !json_encoding_tmp

let () = wrap (Sys.getenv_opt "PPX_ENCODING_MODULE")

let raise_error ~loc s = Location.raise_errorf ~loc s

let enc_names : (string, string) Hashtbl.t = Hashtbl.create 256

let enc_name ?(search=true) type_name =
  let aux name =
    match List.rev @@ String.split_on_char '.' name with
    | [] -> assert false
    | "t" :: q -> String.concat "." @@ List.rev @@ "enc" :: q
    | _ -> name ^ "_enc" in
  if not search then aux type_name
  else match Hashtbl.find_opt enc_names type_name with
    | None -> aux type_name
    | Some name -> name

let add_enc_name type_name enc_name = Hashtbl.add enc_names type_name enc_name

let enc_mod s = !json_encoding_dot ^ s
let enc_var ~loc s = evar ~loc (enc_mod s)
let enc_apply ~loc s l = eapply ~loc (enc_var ~loc s) l
let enc_apply0 ~loc s l = pexp_apply ~loc (enc_var ~loc s) l

let pexp_fun p e = pexp_fun ~loc:e.pexp_loc Nolabel None p e
let llid ~loc s = {txt=Longident.parse s; loc}
let esome e =
  let loc = e.pexp_loc in
  pexp_construct ~loc (llid ~loc "Some") (Some e)
let eoption ~loc = function
  | None -> [%expr None]
  | Some e -> esome e

let rec add_params_fun expr = function
  | [] -> expr
  | {ptyp_desc = Ptyp_var x; ptyp_loc=loc; _} :: t ->
    pexp_fun
      (ppat_constraint ~loc (pvar ~loc ("_" ^ enc_name x))
         (ptyp_constr ~loc (llid ~loc (enc_mod "encoding")) [ptyp_var ~loc x]))
      (add_params_fun expr t)
  | _ :: t -> add_params_fun expr t

let rec add_params_fun_sig typ = function
  | [] -> typ
  | {ptyp_desc = Ptyp_var x; ptyp_loc=loc; _} :: t ->
    ptyp_arrow ~loc Nolabel
      (ptyp_constr ~loc (llid ~loc (enc_mod "encoding")) [ptyp_var ~loc x])
      (add_params_fun_sig typ t)
  | _ :: t -> add_params_fun_sig typ t

let param_names params =
  List.rev @@ List.fold_left (fun acc (p, _) -> match p.ptyp_desc with
    | Ptyp_var x -> x :: acc
    | _ -> acc) [] params

let string_literal = function
  | Ppxlib.Pconst_string (s, _, _) -> Some s
  | _ -> None

let get_expr_attr = function
  | PStr [{pstr_desc = Pstr_eval (e, _); _}] -> Some e
  | _ -> None

let get_string_attr = function
  | PStr [{pstr_desc = Pstr_eval ({pexp_desc = Pexp_constant cst; _}, _); _}] ->
    string_literal cst
  | _ -> None

let rm_prefix_of_expr e = match e.pexp_desc with
  | Pexp_construct ({txt=Lident "true"; _}, _) -> Some (`bool true)
  | Pexp_construct ({txt=Lident "false"; _}, _) -> Some (`bool false)
  | Pexp_constant (Pconst_integer (s, None)) -> Some (`length (int_of_string s))
  | Pexp_constant cst ->
    begin match string_literal cst with
      | Some s -> Some (`prefix s)
      | None -> None
    end
  | _ -> None

let get_rm_prefix_attr pl =
  match get_expr_attr pl with
  | None -> None
  | Some e -> rm_prefix_of_expr e

let to_snake ?(sep='_') s =
  let n = String.length s in
  let b = Bytes.create (2*n) in
  let rec aux i j =
    if i = n then j
    else
    let c = String.get s i in
    let code = Char.code c in
    if i=0 && code >= 65 && code <= 90 then (
      Bytes.set b j (Char.chr (code + 32));
      aux (i+1) (j+1))
    else if code >= 65 && code <= 90 then (
      Bytes.set b j sep;
      Bytes.set b (j+1) (Char.chr (code + 32));
      aux (i+1) (j+2))
    else if c = '-' || c = '_' then (
      Bytes.set b j sep;
      aux (i+1) (j+1))
    else (
      Bytes.set b j c;
      aux (i+1) (j+1)) in
  let m = aux 0 0 in
  Bytes.(to_string @@ sub b 0 m)

let to_camel ?(pascal=false) s =
  let n = String.length s in
  let b = Bytes.create n in
  let rec aux i j =
    if i = n then j
    else
    let c = String.get s i in
    if (c = '_' || c = '-') && i <> (n-1) then (
      let code = Char.code @@ String.get s (i+1) in
      let c =
        if code >= 97 && code <= 122 then Char.chr (code - 32)
        else Char.chr code in
      Bytes.set b j c;
      aux (i+2) (j+1))
    else (
      Bytes.set b j c;
      aux (i+1) (j+1)) in
  let m = aux 0 0 in
  let s = Bytes.(to_string @@ sub b 0 m) in
  if pascal then String.capitalize_ascii s else s

type field_attributes = {
  fa_field : string * bool * expression option;
  fa_key : string;
  fa_title : expression option;
  fa_description : expression option;
  fa_assoc : bool option;
  fa_enum : bool option;
  fa_exclude : expression option;
  fa_obj : bool option;
  fa_enc : expression option;
  fa_obj1 : string option;
  fa_merge : bool;
  fa_construct_default : bool;
  fa_set : expression option;
  fa_map : expression option option;
}

let transform_case ?case ?(typ=`cons) s = match case with
  | None -> begin match typ with `cons -> String.uncapitalize_ascii s | `variant -> s end
  | Some `camel -> to_camel s
  | Some `snake -> to_snake s
  | Some `pascal -> to_camel ~pascal:true s
  | Some `kebab -> to_snake ~sep:'-' s
  | Some `upper -> String.uppercase_ascii s
  | Some `cap -> String.capitalize_ascii s

let field_attrs ~key ?(opt=false) ?(option="opt") ?case ?set ?map l =
  let fa_field = match opt, option with
    | false, _ -> ("req", false, None)
    | _, "opt" -> ("opt", true, None)
    | _, "req" -> ("req", false, None)
    | _ ->
      let loc = !Ast_helper.default_loc in
      ("dft", false, Some [%expr None]) in
  let fa_key = transform_case ?case key in
  List.fold_left (fun fa a -> match a.attr_name.txt with
    | "req" | "enc.req" -> {fa with fa_field = ("req", false, None)}
    | "opt" | "enc.opt" -> {fa with fa_field = ("opt", true, None)}
    | "dft" | "enc.dft" -> {fa with fa_field = ("dft", false, get_expr_attr a.attr_payload)}
    | "ddft" | "enc.ddft" -> {fa with fa_field = ("dft", false, get_expr_attr a.attr_payload);
                                      fa_construct_default = true}
    | "key" | "enc.key" -> begin match get_string_attr a.attr_payload with
      | None -> failwith "key expression must be a string constant"
      | Some fa_key -> {fa with fa_key} end
    | "title" | "enc.title" -> {fa with fa_title = get_expr_attr a.attr_payload}
    | "description" | "enc.description" -> {fa with fa_description = get_expr_attr a.attr_payload}
    | "assoc" | "enc.assoc" -> {fa with fa_assoc = Some true}
    | "enum" | "enc.enum" -> {fa with fa_enum = Some true}
    | "union" | "enc.union" -> {fa with fa_enum = Some false}
    | "exclude" | "enc.exclude" -> {fa with fa_exclude = get_expr_attr a.attr_payload}
    | "object" | "enc.object" -> {fa with fa_obj = Some true}
    | "encoding" | "enc.encoding" -> {fa with fa_enc = get_expr_attr a.attr_payload}
    | "obj1" | "enc.obj1" | "wrap" | "enc.wrap" -> {fa with fa_obj1 = get_string_attr a.attr_payload}
    | "merge" | "enc.merge" -> {fa with fa_merge = true}
    | "camel" | "enc.camel" -> {fa with fa_key = to_camel fa.fa_key}
    | "snake" | "enc.snake" -> {fa with fa_key = to_snake fa.fa_key}
    | "pascal" | "enc.pascal" -> {fa with fa_key = to_camel ~pascal:true fa.fa_key}
    | "kebab" | "enc.kebab" -> {fa with fa_key = to_snake ~sep:'-' fa.fa_key}
    | "set" | "enc.set" -> {fa with fa_set = get_expr_attr a.attr_payload}
    | "map" | "enc.map" -> {fa with fa_map = Some (get_expr_attr a.attr_payload)}
    | "upper" | "enc.upper" -> {fa with fa_key = String.uppercase_ascii fa.fa_key}
    | "cap" | "enc.cap" -> {fa with fa_key = String.capitalize_ascii fa.fa_key}
    | _ -> fa)
    {fa_field; fa_key; fa_title=None; fa_description=None; fa_assoc=None;
     fa_enum=None; fa_exclude=None; fa_obj=None; fa_enc=None; fa_obj1=None;
     fa_merge=false; fa_construct_default=false; fa_set=set; fa_map=map} l

let pp_expression = Pprintast.expression

type cs_attributes = {
  cs_kind : string option;
  cs_kind_label : string option;
  cs_assoc : bool option;
  cs_enum : bool option;
  cs_transform : (string -> string);
  cs_obj : bool option;
  cs_enc : expression option;
  cs_title : expression option;
  cs_description : expression option;
  cs_ignore : bool;
  cs_rm_prefix : [`bool of bool | `prefix of string | `length of int];
  cs_obj1 : string option;
  cs_empty : bool option;
  cs_case : expression option;
}

let constructor_attrs ?case ?typ ?kind l =
  let cs_transform s = transform_case ?case ?typ s in
  let cs_kind, cs_kind_label = match kind with
    | Some (Some k) -> Some "", Some k
    | Some None -> Some "", None
    | _ -> None, None in
  List.fold_left (fun cs a -> match a.attr_name.txt with
    | "kind" | "enc.kind" ->
      let cs_kind = match get_string_attr a.attr_payload with
        | None -> Some "" | k -> k in
      {cs with cs_kind}
    | "kind_label" | "enc.kind_label" -> {cs with cs_kind_label = get_string_attr a.attr_payload}
    | "assoc" | "enc.assoc" -> {cs with cs_assoc = Some true}
    | "enum" | "enc.enum" -> {cs with cs_enum = Some true}
    | "union" | "enc.union" -> {cs with cs_enum = Some false}
    | "key" | "enc.key" -> begin match get_string_attr a.attr_payload with
      | None -> failwith "key expression must be a string constant"
      | Some s -> {cs with cs_transform = fun _ -> s}
    end
    | "object" | "enc.object" -> {cs with cs_obj = Some true}
    | "encoding" | "enc.encoding" -> {cs with cs_enc = get_expr_attr a.attr_payload}
    | "title" | "enc.title" -> {cs with cs_title = get_expr_attr a.attr_payload}
    | "description" | "enc.desccription" -> {cs with cs_description = get_expr_attr a.attr_payload}
    | "ignore" | "enc.ignore" -> {cs with cs_ignore = true}
    | "remove_prefix" | "enc.remove_prefix" ->
      let cs_rm_prefix = match get_rm_prefix_attr a.attr_payload with
        | None -> `bool true
        | Some x -> x in
      {cs with cs_rm_prefix}
    | "obj1" | "enc.obj1" | "wrap" | "enc.wrap" ->
      {cs with cs_obj1 = get_string_attr a.attr_payload}
    | "empty" | "enc.empty" -> {cs with cs_empty = Some true}
    | "camel" | "enc.camel" -> {cs with cs_transform = to_camel}
    | "snake" | "enc.snake" -> {cs with cs_transform = to_snake}
    | "pascal" | "enc.pascal" -> {cs with cs_transform = to_camel ~pascal:true}
    | "kebab" | "enc.kebab" -> {cs with cs_transform = to_snake ~sep:'-'}
    | "upper" | "enc.upper" -> {cs with cs_transform = String.uppercase_ascii}
    | "cap" | "enc.cap" -> {cs with cs_transform = String.capitalize_ascii}
    | "case" | "enc.case" -> {cs with cs_case = get_expr_attr a.attr_payload}
    | _ -> cs)
    { cs_kind; cs_assoc=None; cs_enum=None; cs_transform; cs_obj=None;
      cs_enc=None; cs_title=None; cs_description=None; cs_ignore=false;
      cs_rm_prefix=`bool false; cs_obj1=None; cs_kind_label; cs_empty=None;
      cs_case=None } l

type core_attributes = {
  co_assoc : bool option;
  co_enum : bool option;
  co_exclude : expression option;
  co_obj : bool option;
  co_enc : expression option;
  co_obj1 : string option;
  co_merge : bool;
  co_rm_prefix : [`bool of bool | `prefix of string | `length of int] option;
  co_set: expression option;
  co_map: expression option option;
  co_opt: expression option option;
}

let core_attrs ?assoc ?enum ?obj ?enc ?obj1 ?set ?map l =
  List.fold_left (fun co a -> match a.attr_name.txt with
    | "assoc" | "enc.assoc" -> {co with co_assoc = Some true}
    | "enum" | "enc.enum" -> {co with co_enum = Some true}
    | "union" | "enc.union" -> {co with co_enum = Some false}
    | "exclude" | "enc.exclude" -> {co with co_exclude = get_expr_attr a.attr_payload}
    | "object" | "enc.object" -> {co with co_obj = Some true}
    | "encoding" | "enc.encoding" -> {co with co_enc = get_expr_attr a.attr_payload}
    | "obj1" | "enc.obj1" | "wrap" | "enc.wrap" -> {co with co_obj1 = get_string_attr a.attr_payload}
    | "merge" | "enc.merge" -> {co with co_merge = true}
    | "remove_prefix" | "enc.remove_prefix" -> {co with co_rm_prefix = get_rm_prefix_attr a.attr_payload}
    | "set" -> { co with co_set = get_expr_attr a.attr_payload }
    | "map" -> { co with co_map = Some (get_expr_attr a.attr_payload) }
    | "opt" -> { co with co_opt = Some None }
    | "dft" -> { co with co_opt = Some (get_expr_attr a.attr_payload) }
    | _ -> co)
    {co_assoc=assoc; co_enum=enum; co_exclude=None; co_obj=obj; co_enc=enc;
     co_obj1=obj1; co_merge=false; co_rm_prefix=None; co_set=set; co_map=map;
     co_opt=None} l

let new_var = let i = ref (-1) in fun () -> incr i; "v" ^ string_of_int !i

let str_of_structure e = Pprintast.string_of_structure e
let str_of_signature e =
  Pprintast.signature Format.str_formatter e;
  Format.flush_str_formatter ()

let str_of_core e =
  Pprintast.core_type Format.str_formatter e;
  Format.flush_str_formatter ()

let rec encaps_tuple ~loc var tuple = function
  | [] -> assert false
  | [h] -> var ~loc h
  | h :: t -> tuple ~loc [var ~loc h; encaps_tuple ~loc var tuple t]

let rec encaps_merge ~loc ?(f="merge_objs") = function
  | [] -> assert false
  | [h, merge] -> if merge then h else enc_apply ~loc "obj1" [h]
  | [f1, m1; f2, m2] when not m1 && not m2 -> enc_apply ~loc "obj2" [f1; f2]
  | (h, merge) :: t -> enc_apply ~loc f [
    if merge then h else enc_apply ~loc "obj1" [h];
    encaps_merge ~loc ~f t ]

let obj_expr ~loc ?(kind="obj") l =
  let v = List.mapi (fun i _ -> "x" ^ string_of_int i) l in
  let no_merge = List.for_all (fun (_, b) -> not b) l in
  let n = List.length l in
  if n < 11 && no_merge then
    eapply ~loc (evar ~loc (enc_mod (kind ^ string_of_int n))) (List.map fst l)
  else
  let f = "merge_" ^ kind ^ "s" in
  enc_apply ~loc "conv" [
    pexp_fun (ppat_tuple ~loc (List.map (pvar ~loc) v)) (encaps_tuple ~loc evar pexp_tuple v);
    pexp_fun (encaps_tuple ~loc pvar ppat_tuple v) (pexp_tuple ~loc (List.map (evar ~loc) v));
    encaps_merge ~loc ~f l ]

let remove_prefix s n = String.sub s n (String.length s - n)

let same_prefix l =
  let common_prefix s1 s2 =
    let n1 = String.length s1 in
    let n2 = String.length s2 in
    let rec aux i =
      if i < n1 && i < n2 && s1.[i] = s2.[i] then aux (i+1)
      else i, String.sub s1 0 i in
    aux 0 in
  let rec aux n pr = function
    | [] -> n, pr
    | h :: t ->
      let n, pr = common_prefix h pr in
      aux n pr t in
  let check n l = List.for_all (fun s -> String.length s > n) l in
  match l with
  | [] | [ _ ] -> 0
  | h :: t ->
    let n = fst (aux (String.length h) h t) in
    if check n l then n else 0

let remove_prefix_options l = function
  | Some (`bool false) -> 0
  | Some (`length n) -> n
  | Some (`prefix s) -> String.length s
  | _ -> if not remove_prefix_var then 0 else same_prefix l

module type S = sig
  type 'a encoding
  type 'a field
  type 'a case

  val unit : unit encoding
  val empty : unit encoding
  val int : int encoding
  val int32 : int32 encoding
  val int53 : int64 encoding
  val bool : bool encoding
  val string : string encoding
  val string_enum : (string * 'a) list -> 'a encoding
  val constant : string -> unit encoding
  val bytes : bytes encoding
  val float : float encoding
  val option : 'a encoding -> 'a option encoding

  val req : ?title:string -> ?description:string -> string -> 't encoding -> 't field
  val opt : ?title:string -> ?description:string -> string -> 't encoding -> 't option field
  val dft : ?title:string -> ?description:string -> string -> 't encoding -> 't -> 't field

  val obj1 : 'f1 field -> 'f1 encoding
  val obj2 : 'f1 field -> 'f2 field -> ('f1 * 'f2) encoding
  val obj3 : 'f1 field -> 'f2 field -> 'f3 field -> ('f1 * 'f2 * 'f3) encoding
  val obj4 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> ('f1 * 'f2 * 'f3 * 'f4) encoding
  val obj5 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5) encoding
  val obj6 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6) encoding
  val obj7 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7) encoding
  val obj8 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> 'f8 field -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7 * 'f8) encoding
  val obj9 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> 'f8 field -> 'f9 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7 * 'f8 * 'f9) encoding
  val obj10 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> 'f8 field -> 'f9 field -> 'f10 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7 * 'f8 * 'f9 * 'f10) encoding

  val merge_objs : 'a encoding -> 'b encoding -> ('a * 'b) encoding

  val tup1 : 'f1 field -> 'f1 encoding
  val tup2 : 'f1 field -> 'f2 field -> ('f1 * 'f2) encoding
  val tup3 : 'f1 field -> 'f2 field -> 'f3 field -> ('f1 * 'f2 * 'f3) encoding
  val tup4 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> ('f1 * 'f2 * 'f3 * 'f4) encoding
  val tup5 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5) encoding
  val tup6 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6) encoding
  val tup7 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7) encoding
  val tup8 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> 'f8 field -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7 * 'f8) encoding
  val tup9 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> 'f8 field -> 'f9 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7 * 'f8 * 'f9) encoding
  val tup10 : 'f1 field -> 'f2 field -> 'f3 field -> 'f4 field -> 'f5 field -> 'f6 field
    -> 'f7 field -> 'f8 field -> 'f9 field -> 'f10 field
    -> ('f1 * 'f2 * 'f3 * 'f4 * 'f5 * 'f6 * 'f7 * 'f8 * 'f9 * 'f10) encoding

  val merge_tups : 'a encoding -> 'b encoding -> ('a * 'b) encoding

  val array : 'a encoding -> 'a array encoding
  val list : 'a encoding -> 'a list encoding
  val assoc : 'a encoding -> (string * 'a) list encoding
  val case : ?title:string -> ?description:string -> ('t encoding -> ('a -> 't option) -> ('t -> 'a)) list -> 'a case
  val union : 'a case list -> 'a encoding
  val conv : ('a -> 'b) -> ('b -> 'a) -> 'b encoding -> 'a encoding
  val mu : string -> ?title:string -> ?description:string -> ('a encoding -> 'a encoding) -> 'a encoding
  val any_ezjson_value : 'a encoding
  val def : string -> ?title:string -> ?description:string -> 'a encoding -> 'a encoding
end