open Asttypes
open Parsetree

let expression mapper node =
  let node = Ast_mapper.default_mapper.expr mapper node in
  match node.pexp_desc with
  | Pexp_extension ({ txt = ("continue" | "cps"); loc }, payload) ->
      begin match payload with
      | PStr [{ pstr_desc = Pstr_eval
          ({ pexp_desc = Pexp_let (Nonrecursive, [binding], body); _ }, _); _ }] ->
          let callback = Ast_helper.Exp.fun_ ~loc:binding.pvb_loc
              ~attrs:binding.pvb_attributes Nolabel None binding.pvb_pat body in
          Ast_helper.Exp.apply ~loc:node.pexp_loc ~attrs:node.pexp_attributes
            binding.pvb_expr [Nolabel, callback]
      | _ -> Location.raise_errorf ~loc
          "continuum_ppx: expected let%%continue pattern = expression in body (one nonrecursive binding)"
      end
  | _ -> node

let () = Ast_mapper.register "continuum_ppx"
    (fun _ -> { Ast_mapper.default_mapper with expr = expression })