Source file ppx_custom_printf.ml

open Base
open Ppxlib
open Ast_builder.Default

(* returns the index of the conversion spec (unless the end of string is reached) *)
let rec skip_over_format_flags fmt i =
  if i >= String.length fmt
  then `Eoi
  else (
    match fmt.[i] with
    | '*' | '#' | '-' | ' ' | '+' | '_' | '0' .. '9' | '.' ->
      skip_over_format_flags fmt (i + 1)
    | _ -> `Ok i)
;;

(* doesn't check to make sure the format string is well-formed *)
(* Formats with subformats are skipped for the following reasons:

   One is that they are hard to understand and not often used.

   Another is that subformats like "%(%{Module}%)" won't work, since it
   is impossible to produce a format of type [(Module.t -> 'a,...) format].
*)
let has_subformats (fmt : string) =
  let lim = String.length fmt - 1 in
  let rec loop i =
    if i > lim
    then false
    else if Char.equal fmt.[i] '%'
    then (
      match skip_over_format_flags fmt (i + 1) with
      | `Eoi -> false
      | `Ok i ->
        (match fmt.[i] with
         | '(' | ')' | '}' -> true
         | _ -> loop (i + 1)))
    else loop (i + 1)
  in
  loop 0
;;

(* returns a list of strings where even indexed elements are parts of the format string
   that the preprocessor won't touch and odd indexed elements are the contents of %{...}
   specifications. *)
let explode ~loc (s : string) =
  let len = String.length s in
  (* for cases where we can't parse the string with custom format specifiers, consider
     the string as a regular format string *)
  let as_normal_format_string = [ s ] in
  if has_subformats s
  then as_normal_format_string
  else (
    let sub from to_ = String.sub s ~pos:from ~len:(to_ - from) in
    let rec loop acc from to_ =
      assert (List.length acc % 2 = 0);
      if to_ >= len
      then List.rev (if from >= len then acc else sub from len :: acc)
      else if Char.( <> ) s.[to_] '%'
      then loop acc from (to_ + 1)
      else (
        match skip_over_format_flags s (to_ + 1) with
        | `Eoi -> as_normal_format_string
        | `Ok i ->
          (match s.[i] with
           | '[' ->
             (* Scan char sets are not allowed by printf-like functions. So we might as
                well disallow them at compile-time so that we can reuse them as magic
                format strings in this implementation. *)
             Location.raise_errorf
               ~loc
               "ppx_custom_printf: scan char sets are not allowed in custom format \
                strings"
           | '{' ->
             if to_ + 1 <> i
             then
               Location.raise_errorf
                 ~loc
                 "ppx_custom_printf: unexpected format flags before %%{} specification \
                  in %S"
                 s;
             (match String.index_from s (to_ + 2) '}' with
              | None -> as_normal_format_string
              | Some i ->
                let l = sub (to_ + 2) i :: sub from to_ :: acc in
                loop l (i + 1) (i + 1))
           | _ -> loop acc from (i + 1)))
      (* skip the conversion spec *)
    in
    loop [] 0 0)
;;

let processed_format_string ~exploded_format_string =
  let l =
    let rec loop i l =
      match l with
      | s1 :: _s2 :: l -> s1 :: Printf.sprintf "%%%d[.]" i :: loop (i + 1) l
      | [ s1 ] -> [ s1 ]
      | [] -> []
    in
    loop 0 exploded_format_string
  in
  String.concat l ~sep:""
;;

let rec evens = function
  | ([] | [ _ ]) as l -> l
  | x :: _ :: l -> x :: evens l
;;

let odds = function
  | [] -> []
  | _ :: l -> evens l
;;

(* Returns a pair of:

   - a format string, which is [s] where all custom format specifications have been
     replaced by ["%" ^ string_of_int index ^ "[.]"] where [index] is the number of
     the custom format specification, starting from 0. This string can be passed directly
     to [CamlinternalFormat.fmt_ebb_of_string]
   - an array of custom format specifications, in the order they appear in the original
     string
*)
let extract_custom_format_specifications ~loc s =
  let exploded_format_string = explode ~loc s in
  let processed = processed_format_string ~exploded_format_string in
  let custom_specs = Array.of_list (odds exploded_format_string) in
  processed, custom_specs
;;

let gen_symbol = gen_symbol ~prefix:"_custom_printf"

let is_space = function
  | ' ' | '\t' | '\n' | '\r' -> true
  | _ -> false
;;

let strip s =
  let a = ref 0 in
  let b = ref (String.length s - 1) in
  while !a <= !b && is_space s.[!a] do
    Int.incr a
  done;
  while !a <= !b && is_space s.[!b] do
    Int.decr b
  done;
  if !a > !b then "" else String.sub s ~pos:!a ~len:(!b - !a + 1)
;;

let string_to_expr ~loc s =
  let sexp_converter_opt =
    match String.lsplit2 s ~on:':' with
    | None -> None
    | Some ("sexp", colon_suffix) ->
      Some ([%expr Ppx_sexp_conv_lib.Sexp.to_string_hum], colon_suffix)
    | Some (colon_prefix, colon_suffix) ->
      (match String.chop_prefix colon_prefix ~prefix:"sexp#" with
       | None -> None
       | Some hash_suffix ->
         Some
           ( pexp_ident
               ~loc
               (Located.mk
                  ~loc
                  (Longident.parse ("Ppx_sexp_conv_lib.Sexp.to_string_" ^ hash_suffix)))
           , colon_suffix ))
  in
  match sexp_converter_opt with
  | Some (sexp_converter, unparsed_type) ->
    let lexbuf = Lexing.from_string unparsed_type in
    (* ~loc is the position of the string, not the position of the %{bla} group we're
       looking at. The format strings don't contain location information, so we can't
       actually find the proper positions. *)
    lexbuf.lex_abs_pos <- loc.loc_start.pos_cnum;
    lexbuf.lex_curr_p <- loc.loc_start;
    let ty = Parse.core_type lexbuf in
    let e = Ppx_sexp_conv_expander.Sexp_of.core_type ty in
    let arg = gen_symbol () in
    pexp_fun
      ~loc
      Nolabel
      None
      (pvar ~loc arg)
      (eapply ~loc sexp_converter [ eapply ~loc e [ evar ~loc arg ] ])
  | None ->
    let fail loc =
      Location.raise_errorf
        ~loc
        "ppx_custom_printf: string %S should be of the form <Module>, \
         <Module>.<identifier>, <Module>#identifier, sexp:<type>, or sexp#mach:<type>"
        s
    in
    let s, has_hash_suffix, to_string =
      match String.lsplit2 s ~on:'#' with
      | None -> s, false, "to_string"
      | Some (s, hash_suffix) -> s, true, "to_string_" ^ hash_suffix
    in
    let to_string_id : Longident.t =
      let s = strip s in
      match s with
      | "" -> Lident to_string
      | _ ->
        (match Longident.parse s with
         | (Lident n | Ldot (_, n)) as id ->
           if String.( <> ) n "" && Char.equal (Char.uppercase n.[0]) n.[0]
           then Longident.Ldot (id, to_string)
           else if not has_hash_suffix
           then id
           else fail loc
         | _ -> fail loc)
    in
    let func = pexp_ident ~loc (Located.mk ~loc to_string_id) in
    (* Eta-expand as the to_string function might take optional arguments *)
    let arg = gen_symbol () in
    pexp_fun ~loc Nolabel None (pvar ~loc arg) (eapply ~loc func [ evar ~loc arg ])
;;

class lifter ~loc ~custom_specs =
  object (self)
    inherit [expression] Format_lifter.lift as super
    inherit Ppxlib_metaquot_lifters.expression_lifters loc

    method! fmt
      : type f0
      f1
      f2
      f3
      f4
      f5.  (f0 -> expression)
          -> (f1 -> expression)
          -> (f2 -> expression)
          -> (f3 -> expression)
          -> (f4 -> expression)
          -> (f5 -> expression)
          -> (f0, f1, f2, f3, f4, f5) CamlinternalFormatBasics.fmt
          -> expression =
      fun f0 f1 f2 f3 f4 f5 fmt ->
        let open CamlinternalFormatBasics in
        match fmt with
        (* Recognize the special form "%index[...whatever...]" *)
        | Scan_char_set (Some idx, _, fmt)
        (* [custom_specs] is empty if [explode] couldn't parse the string. In this case we
           can have some scar char sets left. *)
          when idx >= 0 && idx < Array.length custom_specs ->
          let rest = self#fmt (fun _ -> assert false) f1 f2 f3 f4 f5 fmt in
          let func = string_to_expr ~loc custom_specs.(idx) in
          [%expr Custom (Custom_succ Custom_zero, (fun () -> [%e func]), [%e rest])]
        | _ -> super#fmt f0 f1 f2 f3 f4 f5 fmt
  end

let expand_format_string ~loc fmt_string =
  let processed_fmt_string, custom_specs =
    extract_custom_format_specifications ~loc fmt_string
  in
  let (CamlinternalFormat.Fmt_EBB fmt) =
    try CamlinternalFormat.fmt_ebb_of_string processed_fmt_string with
    | e ->
      Location.raise_errorf
        ~loc
        "%s"
        (match e with
         (* [fmt_ebb_of_string] normally raises [Failure] on invalid input *)
         | Failure msg -> msg
         | e -> Exn.to_string e)
  in
  let lifter = new lifter ~loc ~custom_specs in
  let format6 = CamlinternalFormatBasics.Format (fmt, fmt_string) in
  let phantom _ = assert false in
  let e = lifter#format6 phantom phantom phantom phantom phantom phantom format6 in
  [%expr ([%e e] : (_, _, _, _, _, _) CamlinternalFormatBasics.format6)]
;;

let expand e =
  match e.pexp_desc with
  | Pexp_apply
      ( { pexp_attributes = ident_attrs; _ }
      , [ ( Nolabel
          , { pexp_desc = Pexp_constant (Pconst_string (str, _, _))
            ; pexp_loc = loc
            ; pexp_loc_stack = _
            ; pexp_attributes = str_attrs
            } )
        ] ) ->
    assert_no_attributes ident_attrs;
    assert_no_attributes str_attrs;
    let e' = expand_format_string ~loc str in
    Some { e' with pexp_attributes = Merlin_helpers.hide_attribute :: e.pexp_attributes }
  | _ -> None
;;

let () =
  Driver.register_transformation
    "custom_printf"
    ~rules:[ Context_free.Rule.special_function "!" expand ]
;;