Source file `ppx_pipebang.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
open! Ppxlib

let expand (e : Parsetree.expression) =
  match e.pexp_desc with
  | Pexp_apply (_, [(Nolabel, x); (Nolabel, y)]) ->
      Some (
        match y with
        | { pexp_desc = Pexp_construct (id, None); _ } ->
          { y with pexp_desc = Pexp_construct (id, Some x) }
        | { pexp_desc = Pexp_apply (f, args); pexp_attributes = []; _ }
          when (match f.pexp_desc with
            (* Do not inline |> as this would create applications with too many
               arguments *)
            | Pexp_ident { txt = Lident "|>"; _ } -> false
            | _ -> true) ->
          { e with pexp_desc = Pexp_apply (f, args @ [(Nolabel, x)]) }
        | _ ->
          { e with pexp_desc = Pexp_apply (y, [(Nolabel, x)]) }
      )
  | Pexp_ident { txt = Lident s; _ }
  | Pexp_apply ({ pexp_desc = Pexp_ident { txt = Lident s; _ }; _ }, _) ->
    Location.raise_errorf ~loc:e.pexp_loc "%s must be applied to two arguments" s
  | _ -> None
;;

let () =
  Driver.register_transformation "pipebang"
    ~rules:[ Context_free.Rule.special_function "|>" expand ]
;;