ppx_jsx_tyxml.ml

open Ppxlib

module B = Ast_builder.Default


let contains_jsx : attributes -> bool =
  List.exists (function
    | { txt = "JSX"; _ }, _ -> true
    | _ -> false
  )


let ocaml_reserved =
  [ "and"; "as"; "asr"; "assert"; "begin"; "class"; "constraint"; "do"; "done"
  ; "downto"; "else"; "end"; "exception"; "external"; "false"; "for"; "fun"
  ; "function"; "functor"; "if"; "in"; "include"; "inherit"; "initializer"
  ; "land"; "lazy"; "let"; "lor"; "lsl"; "lsr"; "lxor"; "match"; "method"
  ; "mod"; "module"; "open"; "mutable"; "new"; "nonrec"; "object"; "of"; "open"
  ; "open!"; "or"; "private"; "rec"; "sig"; "struct"; "then"; "to"; "true"
  ; "try"; "type"; "val"; "virtual"; "when"; "while"; "with" ]


let transform_attr_name name =
  let len = String.length name in
  let last = name.[len - 1] in
  let name_trim = String.sub name 0 (len - 1) in
  if last = '_' && List.mem name_trim ocaml_reserved then
    "a_" ^ name_trim
  else
    "a_" ^ name


let rec map_tail ~f expr = match expr.pexp_desc with
  | Pexp_construct ({ txt = Lident "::"; _ } as cons, Some args) ->
    let args = map_args ~f args in
    { expr with pexp_desc = Pexp_construct (cons, Some args) }
  | Pexp_construct ({ txt = Lident "[]"; _ }, None) ->
    expr
  | _ ->
    Location.raise_errorf ~loc:expr.pexp_loc
      "map_elist: expected list constructor or tail"

and map_args ~f expr = match expr.pexp_desc with
  | Pexp_tuple [head; tail] ->
    let head = f head in
    let tail = map_tail ~f tail in
    { expr with pexp_desc = Pexp_tuple [head; tail] }
  | Pexp_tuple _ ->
    Location.raise_errorf ~loc:expr.pexp_loc
      "map_elist: expected 2-tuple, got n-tuple"
  | _ ->
    Location.raise_errorf ~loc:expr.pexp_loc
      "map_elist: expected 2-tuple"


let transform_consts expr = match expr.pexp_desc with
  | Pexp_constant (Pconst_string _) ->
    let loc = expr.pexp_loc in
    [%expr txt [%e expr]]
  | _ -> expr


let transform_args args =
  let rec f acc = function
    | (Labelled "children", children) :: (Nolabel, [%expr ()]) :: [] ->
      let args = List.rev acc |> B.elist ~loc:children.pexp_loc in
      let children = map_tail ~f:transform_consts children in
      [(Labelled "a", args); (Nolabel, children)]

    | (Labelled name, expr) :: rest ->
      let arg =
        B.(pexp_apply ~loc:expr.pexp_loc
          (pexp_ident ~loc:expr.pexp_loc
            (Located.lident ~loc:expr.pexp_loc (transform_attr_name name)))
          [Nolabel, expr]
        )
      in
      f (arg :: acc) rest

    | _ ->
      invalid_arg "expected labeled arg, got unlabeled"
  in
  f [] args


let transform_expr = function
  | Pexp_apply (
      { pexp_desc = Pexp_ident
        { txt = Ldot (lident, "createElement"); loc = lident_loc}
      ; pexp_loc = ident_loc; _ },
      args
    ) ->
    Pexp_apply (
      B.(pexp_ident ~loc:ident_loc
        (Located.mk ~loc:lident_loc (Ldot (lident, "view")))),
      transform_args args
    )

  | Pexp_apply (tag, args) ->
    Pexp_apply (
      tag,
      transform_args args
    )

  | Pexp_construct ({ txt = Lident "::"; _ } as cons, Some children) ->
    let children = map_args ~f:transform_consts children in
    Pexp_construct (cons, Some children)

  | expr ->
    expr


let map_exprs = object
  inherit Ast_traverse.map as super

  method! expression expr =
    let expr = super#expression expr in
    if contains_jsx expr.pexp_attributes then
      { expr with pexp_desc = transform_expr expr.pexp_desc }
    else
      expr
end


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