Source file ppx_cold.ml

open Base
open Ppxlib

let payload_never ~loc =
  let (module B) = Ast_builder.make loc in
  let open B in
  PStr [ pstr_eval (pexp_ident (Located.lident "never")) [] ]
;;

let expand_cold_attribute attr =
  assert (String.equal attr.attr_name.txt "cold");
  let loc = { attr.attr_name.loc with loc_ghost = true } in
  let payload = payload_never ~loc in
  [ Loc.make ~loc "ocaml.inline", payload
  ; Loc.make ~loc "ocaml.local", payload
  ; Loc.make ~loc "ocaml.specialise", payload
  ]
  |> List.map ~f:(fun (name, payload) ->
       Ast_builder.Default.attribute ~loc ~name ~payload)
;;

class attributes_mapper =
  object
    inherit Ast_traverse.map as super

    method! attributes attrs =
      let attrs =
        List.concat_map attrs ~f:(function
          | { attr_name = { txt = "cold"; _ }; attr_payload = PStr []; attr_loc = _ } as
            attr ->
            Attribute.mark_as_handled_manually attr;
            expand_cold_attribute attr
          | attr -> [ attr ])
      in
      super#attributes attrs
  end

let expand_cold = (new attributes_mapper)#structure
let () = Driver.register_transformation "cold" ~impl:expand_cold