Source file ppx.ml

open Ppxlib
module Builder = Ast_builder.Default

module Private = struct
  module Typemod_hide = struct
    let no_type_defined (x : structure_item) =
      match x.pstr_desc with
      | Pstr_eval _ | Pstr_value _ | Pstr_primitive _ | Pstr_typext _ | Pstr_exception _
      (* | Pstr_module {pmb_expr = {pmod_desc = Pmod_ident _} }  *) ->
          true
      | Pstr_include
          {
            pincl_mod =
              {
                pmod_desc =
                  Pmod_constraint ({ pmod_desc = Pmod_structure [ { pstr_desc = Pstr_primitive _; _ } ]; _ }, _);
                _;
              };
            _;
          } ->
          true
          (* FIX https://github.com/rescript-lang/rescript/issues/4881
             generated code from:
             {[
               external %private x : int -> int = "x"
               [@@mel.module "./x"]
             ]}
          *)
      | _ -> false

    let check (x : structure) =
      List.iter
        (fun x ->
          if not (no_type_defined x) then
            Location.raise_errorf ~loc:x.pstr_loc "the structure is not supported in local extension")
        x
  end

  let rule =
    let expand (stru : structure) =
      Typemod_hide.check stru;
      let last_loc = (List.hd stru).pstr_loc in
      let first_loc = (List.hd stru).pstr_loc in
      let loc = { first_loc with loc_end = last_loc.loc_end } in
      Ast_helper.[ Str.open_ (Opn.mk ~override:Override (Mod.structure ~loc stru)) ] |> List.hd
    in
    let rule label =
      let extractor = Ast_pattern.__' in
      let handler ~ctxt:_ { txt = payload; loc } =
        match payload with
        | PStr work -> expand work
        | PSig _ | PTyp _ | PPat _ -> Location.raise_errorf ~loc "private extension is not support"
      in

      let extender = Extension.V3.declare label Structure_item extractor handler in
      Context_free.Rule.extension extender
    in
    rule "private"
end

module Mel_module = struct
  type bundler = Webpack | Esbuild

  let bundler = ref Webpack
  let prefix = ref "/"
  let is_melange_attr { attr_name = { txt = attr } } = "mel.module" = attr
  let has_attr attrs = List.exists is_melange_attr attrs

  let asset_payload attrs =
    let attr =
      (* we use `find` directly even if it can raise, assuming `has_attr` has been called before *)
      List.find is_melange_attr attrs
    in
    match attr.attr_payload with
    | PStr [ { pstr_desc = Pstr_eval ({ pexp_desc = Pexp_constant (Pconst_string (str, _, _)) }, _) } ]
      when String.length (Filename.extension str) > 0 ->
        Some str
    | _ -> None

  module Esbuild = struct
    (* This code is adapted from Esbuild hashing algorithm:
       base32: https://github.com/evanw/esbuild/blob/efa3dd2d8e895f7f9a9bef0d588560bbae7d776e/internal/bundler/bundler.go#L1174
       sum function: https://github.com/evanw/esbuild/blob/efa3dd2d8e895f7f9a9bef0d588560bbae7d776e/internal/xxhash/xxhash.go#L104
       the internal xxhash that esbuild uses is adapted from https://github.com/cespare/xxhash
    *)
    let hash_for_filename bytes = String.sub (Base32.encode_string (Bytes.to_string bytes)) 0 8

    let sum hex_str =
      (* Convert hexadecimal string to Int64 *)
      let int64_value = Int64.of_string ("0x" ^ hex_str) in

      (* Create an 8-byte buffer *)
      let bytes = Bytes.create 8 in

      (* Fill the buffer with the bytes of the Int64 value *)
      for i = 0 to 7 do
        let byte = Int64.(to_int (shift_right_logical int64_value (8 * (7 - i)))) land 0xFF in
        Bytes.set bytes i (char_of_int byte)
      done;

      bytes

    let hash content =
      let hash = XXH64.hash content in
      let b = sum (XXH64.to_hex hash) in
      hash_for_filename b

    let filename ~base content = Filename.(chop_extension base ^ "-" ^ hash content ^ extension base)
  end

  (*
     (* For now, rspack doesn't support real content hashes, see https://github.com/web-infra-dev/rspack/issues/6606 *)
      module Rspack = struct
         (* This code is adapted from Rspack hashing algorithm:
            https://github.com/web-infra-dev/rspack/blob/0a5cf0ddf38d41c2cad58c95ee9c1d3bd95e377f/crates/rspack_hash/src/lib.rs
         *)
         let hex_to_little_endian hex_str =
           (* Split the hex string into byte pairs *)
           let rec split_into_bytes acc i =
             if i >= String.length hex_str then List.rev acc
             else
               let byte = String.sub hex_str i 2 in
               split_into_bytes (byte :: acc) (i + 2)
           in
           (* Join byte pairs into a single string *)
           let join_bytes bytes = String.concat "" bytes in
           (* Perform the transformation *)
           let bytes = split_into_bytes [] 0 in
           let reversed_bytes = List.rev bytes in
           join_bytes reversed_bytes

         let hash content =
           let open XXHash in
           let hash = XXH3_64.hash content in
           hex_to_little_endian (XXH3_64.to_hex hash)
       end *)
  module Webpack = struct
    (* Needs following config in webpack.config.js, see https://webpack.js.org/configuration/output/#outputhashfunction
       ```
       module.exports = {
         //...
         output: {
           hashFunction: 'xxhash64',
         },
       };
       ```
       Also needs to set `realContentHash` for it to work in dev mode (see https://webpack.js.org/configuration/optimization/#optimizationrealcontenthash):
       ```
       module.exports = {
         //...
         optimization: {
           realContentHash: false,
         },
       };
       ```
    *)
    let hash content =
      let hash = XXH64.hash content in
      XXH64.to_hex hash

    let filename ~base content = hash content ^ Filename.extension base
  end
end

module String_interpolation = struct
  (* https://github.com/melange-re/melange/blob/fb1466fed7d6e5aafd3ee266bbd4ec70c8fb857a/ppx/string_interp.ml *)
  module Utf8_string = struct
    type byte = Single of int | Cont of int | Leading of int * int | Invalid

    (** [classify chr] returns the {!byte} corresponding to [chr] *)
    let classify chr =
      let c = int_of_char chr in
      (* Classify byte according to leftmost 0 bit *)
      if c land 0b1000_0000 = 0 then Single c
      else if
        (* c 0b0____*)
        c land 0b0100_0000 = 0
      then Cont (c land 0b0011_1111)
      else if
        (* c 0b10___*)
        c land 0b0010_0000 = 0
      then Leading (1, c land 0b0001_1111)
      else if
        (* c 0b110__*)
        c land 0b0001_0000 = 0
      then Leading (2, c land 0b0000_1111)
      else if
        (* c 0b1110_ *)
        c land 0b0000_1000 = 0
      then Leading (3, c land 0b0000_0111)
      else if
        (* c 0b1111_0___*)
        c land 0b0000_0100 = 0
      then Leading (4, c land 0b0000_0011)
      else if
        (* c 0b1111_10__*)
        c land 0b0000_0010 = 0
      then Leading (5, c land 0b0000_0001) (* c 0b1111_110__ *)
      else Invalid
  end

  type error =
    | Invalid_code_point
    | Unterminated_backslash
    | Unterminated_variable
    | Unmatched_paren
    | Invalid_syntax_of_var of string

  type kind = String | Var of int * int
  (* [Var (loffset, roffset)]
     For parens it used to be (2,-1)
     for non-parens it used to be (1,0) *)

  (* Note the position is about code point *)
  type pos = {
    lnum : int;
    offset : int;
    byte_bol : int; (* Note it actually needs to be in sync with OCaml's lexing semantics *)
  }

  type segment = { start : pos; finish : pos; kind : kind; content : string }

  type cxt = {
    mutable segment_start : pos;
    buf : Buffer.t;
    s_len : int;
    mutable segments : segment list;
    pos_bol : int; (* record the abs position of current beginning line *)
    byte_bol : int;
    pos_lnum : int; (* record the line number *)
  }

  exception Error of pos * pos * error

  let pp_error fmt err =
    Format.pp_print_string fmt
    @@
    match err with
    | Invalid_code_point -> "Invalid code point"
    | Unterminated_backslash -> "\\ ended unexpectedly"
    | Unterminated_variable -> "$ unterminated"
    | Unmatched_paren -> "Unmatched paren"
    | Invalid_syntax_of_var s -> "`" ^ s ^ "' is not a valid syntax of interpolated identifer"

  let valid_lead_identifier_char x = match x with 'a' .. 'z' | '_' -> true | _ -> false
  let valid_identifier_char x = match x with 'a' .. 'z' | 'A' .. 'Z' | '0' .. '9' | '_' | '\'' -> true | _ -> false

  (* Invariant: [valid_lead_identifier] has to be [valid_identifier] *)
  let valid_identifier =
    let for_all_from =
      let rec unsafe_for_all_range s ~start ~finish p =
        start > finish || (p (String.unsafe_get s start) && unsafe_for_all_range s ~start:(start + 1) ~finish p)
      in
      fun s start p ->
        let len = String.length s in
        if start < 0 then invalid_arg "for_all_from" else unsafe_for_all_range s ~start ~finish:(len - 1) p
    in
    fun s ->
      let s_len = String.length s in
      if s_len = 0 then false else valid_lead_identifier_char s.[0] && for_all_from s 1 valid_identifier_char

  (* FIXME: multiple line offset
     if there is no line offset. Note {|{j||} border will never trigger a new
     line *)
  let update_position border { lnum; offset; byte_bol } (pos : Lexing.position) =
    if lnum = 0 then { pos with pos_cnum = pos.pos_cnum + border + offset }
      (* When no newline, the column number is [border + offset] *)
    else
      {
        pos with
        pos_lnum = pos.pos_lnum + lnum;
        pos_bol = pos.pos_cnum + border + byte_bol;
        pos_cnum = pos.pos_cnum + border + byte_bol + offset;
        (* when newline, the column number is [offset] *)
      }

  let update border start finish (loc : Location.t) =
    let start_pos = loc.loc_start in
    { loc with loc_start = update_position border start start_pos; loc_end = update_position border finish start_pos }

  let pos_error cxt ~loc error =
    raise
      (Error (cxt.segment_start, { lnum = cxt.pos_lnum; offset = loc - cxt.pos_bol; byte_bol = cxt.byte_bol }, error))

  let add_var_segment cxt loc loffset roffset =
    let content = Buffer.contents cxt.buf in
    Buffer.clear cxt.buf;
    let next_loc = { lnum = cxt.pos_lnum; offset = loc - cxt.pos_bol; byte_bol = cxt.byte_bol } in
    if valid_identifier content then (
      cxt.segments <-
        { start = cxt.segment_start; finish = next_loc; kind = Var (loffset, roffset); content } :: cxt.segments;
      cxt.segment_start <- next_loc)
    else
      let cxt =
        match String.trim content with
        | "" ->
            (* Move the position back 2 characters "$(" if this is the empty
               interpolation. *)
            {
              cxt with
              segment_start =
                {
                  cxt.segment_start with
                  offset = (match cxt.segment_start.offset with 0 -> 0 | n -> n - 3);
                  byte_bol = (match cxt.segment_start.byte_bol with 0 -> 0 | n -> n - 3);
                };
              pos_bol = cxt.pos_bol + 3;
              byte_bol = cxt.byte_bol + 3;
            }
        | _ -> cxt
      in
      pos_error cxt ~loc (Invalid_syntax_of_var content)

  let add_str_segment cxt loc =
    let content = Buffer.contents cxt.buf in
    Buffer.clear cxt.buf;
    let next_loc = { lnum = cxt.pos_lnum; offset = loc - cxt.pos_bol; byte_bol = cxt.byte_bol } in
    cxt.segments <- { start = cxt.segment_start; finish = next_loc; kind = String; content } :: cxt.segments;
    cxt.segment_start <- next_loc

  let rec check_and_transform loc s byte_offset ({ s_len; buf; _ } as cxt) =
    if byte_offset = s_len then add_str_segment cxt loc
    else
      let current_char = s.[byte_offset] in
      match Utf8_string.classify current_char with
      | Single 92 (* '\\' *) ->
          let loc = loc + 1 in
          let offset = byte_offset + 1 in
          if offset >= s_len then pos_error cxt ~loc Unterminated_backslash else Buffer.add_char buf '\\';
          let cur_char = s.[offset] in
          Buffer.add_char buf cur_char;
          check_and_transform (loc + 1) s (offset + 1) cxt
      | Single 36 ->
          (* $ *)
          add_str_segment cxt loc;
          let offset = byte_offset + 1 in
          if offset >= s_len then pos_error ~loc cxt Unterminated_variable
          else
            let cur_char = s.[offset] in
            if cur_char = '(' then expect_var_paren (loc + 2) s (offset + 1) cxt
            else expect_simple_var (loc + 1) s offset cxt
      | Single _ | Leading _ | Cont _ ->
          Buffer.add_char buf current_char;
          check_and_transform (loc + 1) s (byte_offset + 1) cxt
      | Invalid -> pos_error ~loc cxt Invalid_code_point

  (* Lets keep identifier simple, so that we could generating a function easier
     in the future for example
     let f = [%fn{| $x + $y = $x_add_y |}] *)
  and expect_simple_var loc s offset ({ buf; s_len; _ } as cxt) =
    let v = ref offset in
    if not (offset < s_len && valid_lead_identifier_char s.[offset]) then
      pos_error cxt ~loc (Invalid_syntax_of_var String.empty)
    else (
      while !v < s_len && valid_identifier_char s.[!v] do
        (* TODO *)
        let cur_char = s.[!v] in
        Buffer.add_char buf cur_char;
        incr v
      done;
      let added_length = !v - offset in
      let loc = added_length + loc in
      add_var_segment cxt loc 1 0;
      check_and_transform loc s (added_length + offset) cxt)

  and expect_var_paren loc s offset ({ buf; s_len; _ } as cxt) =
    let v = ref offset in
    while !v < s_len && s.[!v] <> ')' do
      let cur_char = s.[!v] in
      Buffer.add_char buf cur_char;
      incr v
    done;
    let added_length = !v - offset in
    let loc = added_length + 1 + loc in
    if !v < s_len && s.[!v] = ')' then (
      add_var_segment cxt loc 2 (-1);
      check_and_transform loc s (added_length + 1 + offset) cxt)
    else pos_error cxt ~loc Unmatched_paren

  (* TODO: Allow identifers x.A.y *)

  let border = String.length "{j|"

  let rec handle_segments =
    let module Exp = Ast_helper.Exp in
    let concat_ident : Longident.t = Ldot (Lident "Stdlib", "^") in
    let escaped_js_delimiter =
      (* syntax not allowed at the user level *)
      let unescaped_js_delimiter = "js" in
      Some unescaped_js_delimiter
    in
    let merge_loc (l : Location.t) (r : Location.t) =
      if l.loc_ghost then r
      else if r.loc_ghost then l
      else
        match (l, r) with
        | { loc_start; _ }, { loc_end; _ } (* TODO: improve*) -> { loc_start; loc_end; loc_ghost = false }
    in
    let aux loc segment =
      match segment with
      | { start; finish; kind; content } -> (
          match kind with
          | String ->
              let loc = update border start finish loc in
              Exp.constant (Pconst_string (content, loc, escaped_js_delimiter))
          | Var (soffset, foffset) ->
              let loc =
                {
                  loc with
                  loc_start = update_position (soffset + border) start loc.loc_start;
                  loc_end = update_position (foffset + border) finish loc.loc_start;
                }
              in
              Exp.ident ~loc { loc; txt = Lident content })
    in
    let concat_exp a_loc x ~(lhs : expression) =
      let loc = merge_loc a_loc lhs.pexp_loc in
      Exp.apply (Exp.ident { txt = concat_ident; loc }) [ (Nolabel, lhs); (Nolabel, aux loc x) ]
    in
    fun loc rev_segments ->
      match rev_segments with
      | [] -> Exp.constant (Pconst_string ("", loc, escaped_js_delimiter))
      | [ segment ] -> aux loc segment (* string literal *)
      | { content = ""; _ } :: rest -> handle_segments loc rest
      | a :: rest -> concat_exp loc a ~lhs:(handle_segments loc rest)

  let transform =
    let transform (e : expression) s =
      let s_len = String.length s in
      let buf = Buffer.create (s_len * 2) in
      let cxt =
        {
          segment_start = { lnum = 0; offset = 0; byte_bol = 0 };
          buf;
          s_len;
          segments = [];
          pos_lnum = 0;
          byte_bol = 0;
          pos_bol = 0;
        }
      in
      check_and_transform 0 s 0 cxt;
      handle_segments e.pexp_loc cxt.segments
    in
    fun ~loc expr s ->
      try transform expr s
      with Error (start, pos, error) ->
        let loc = update border start pos loc in
        Location.raise_errorf ~loc "%a" pp_error error
end

let is_send_pipe pval_attributes =
  List.exists (fun { attr_name = { txt = attr } } -> String.equal attr "mel.send.pipe") pval_attributes

let has_browser_ppx_attribute attrs =
  List.exists
    (fun { attr_name = { txt = attr }; attr_payload; _ } ->
      match (attr, attr_payload) with
      | "browser_only", _ -> true
      | "platform", PStr [ { pstr_desc = Pstr_eval ({ pexp_desc = Pexp_ident { txt = Lident "js" } }, _); _ } ] -> true
      | _ -> false)
    attrs

let has_attribute attrs name = List.exists (fun { attr_name = { txt; _ }; _ } -> String.equal txt name) attrs

(* Keep this keyword list and translate_mel_obj_label in sync with Melange's
   Lam_methname.translate implementation:
   https://github.com/melange-re/melange/blob/main/jscomp/common/lam_methname.ml *)
let mel_obj_keywords =
  [
    "and";
    "as";
    "assert";
    "begin";
    "class";
    "constraint";
    "do";
    "done";
    "downto";
    "else";
    "end";
    "exception";
    "external";
    "false";
    "for";
    "fun";
    "function";
    "functor";
    "if";
    "in";
    "include";
    "inherit";
    "initializer";
    "lazy";
    "let";
    "match";
    "method";
    "module";
    "mutable";
    "new";
    "nonrec";
    "object";
    "of";
    "open";
    "or";
    "private";
    "rec";
    "sig";
    "struct";
    "then";
    "to";
    "true";
    "try";
    "type";
    "val";
    "virtual";
    "when";
    "while";
    "with";
    "mod";
    "land";
    "lor";
    "lxor";
    "lsl";
    "lsr";
    "asr";
  ]

let find_double_underscore name =
  let rec go index =
    if index < 0 then -1
    else if index + 1 < String.length name && name.[index] = '_' && name.[index + 1] = '_' then index
    else go (index - 1)
  in
  go (String.length name - 2)

let translate_mel_obj_label name =
  let valid_start_char = function '_' | 'a' .. 'z' -> true | _ -> false in
  let double_underscore_index = find_double_underscore name in
  if double_underscore_index = 0 then name
  else if double_underscore_index > 0 then String.sub name 0 double_underscore_index
  else
    match name.[0] with
    | '_' when String.length name > 1 ->
        let candidate = String.sub name 1 (String.length name - 1) in
        if (not (valid_start_char candidate.[0])) || List.mem candidate mel_obj_keywords then candidate else name
    | _ -> name

type js_object_field = { method_name : string; js_name : string; present_expr : expression; value_expr : expression }

let option_is_some_expr ~loc expr = [%expr match [%e expr] with None -> false | Some _ -> true]

let js_obj_internal_expression ~loc name =
  Builder.pexp_ident ~loc { txt = Ldot (Ldot (Ldot (Lident "Js", "Obj"), "Internal"), name); loc }

let is_literal_bool_expr expr =
  match expr.pexp_desc with Pexp_construct ({ txt = Lident ("true" | "false"); _ }, None) -> true | _ -> false

(* Emits a registered [Js.t] object. Field values live in plain refs read by
   the object methods; the registry entry is a single [Deferred] thunk that
   builds the per-field [entry] records (a record plus two boxing closures
   each) only if [Js.Obj.keys]/[assign]/[merge] ever inspects the object.
   This keeps object construction — every [makeProps] call and [%mel.obj]
   literal — free of per-field entry allocation. *)
let build_registered_js_object_expression ?as_type ?(register_name = "register_deferred") ~loc fields =
  let generated_fields =
    List.mapi
      (fun index { method_name; js_name; present_expr; value_expr } ->
        let cell_name = Printf.sprintf "__js_obj_cell_%d" index in
        let cell_binding =
          Builder.value_binding ~loc
            ~pat:(Builder.ppat_var ~loc { loc; txt = cell_name })
            ~expr:[%expr Stdlib.ref [%e value_expr]]
        in
        (* [present] must be evaluated at construction time (it reads the
           original argument), but literal booleans can be inlined into the
           thunk directly. *)
        let present_bindings, present_in_thunk =
          if is_literal_bool_expr present_expr then ([], present_expr)
          else
            let present_name = Printf.sprintf "__js_obj_present_%d" index in
            ( [ Builder.value_binding ~loc ~pat:(Builder.ppat_var ~loc { loc; txt = present_name }) ~expr:present_expr ],
              Builder.evar ~loc present_name )
        in
        let entry_expr =
          Builder.pexp_apply ~loc
            (js_obj_internal_expression ~loc "deferred_entry")
            [
              (Labelled "method_name", Builder.estring ~loc method_name);
              (Labelled "js_name", Builder.estring ~loc js_name);
              (Labelled "present", present_in_thunk);
              (Nolabel, Builder.evar ~loc cell_name);
            ]
        in
        let method_body = Builder.pexp_apply ~loc (Builder.evar ~loc "!") [ (Nolabel, Builder.evar ~loc cell_name) ] in
        let method_ =
          Builder.pcf_method ~loc (Builder.Located.mk method_name ~loc, Public, Cfk_concrete (Fresh, method_body))
        in
        (cell_binding :: present_bindings, entry_expr, method_))
      fields
  in
  let field_bindings, entry_exprs, methods =
    List.fold_right
      (fun (bindings, entry_expr, method_) (field_bindings, entry_exprs, methods) ->
        (bindings @ field_bindings, entry_expr :: entry_exprs, method_ :: methods))
      generated_fields ([], [], [])
  in
  let object_name = "__js_obj" in
  let object_binding =
    Builder.value_binding ~loc
      ~pat:(Builder.ppat_var ~loc { loc; txt = object_name })
      ~expr:(Builder.pexp_object ~loc (Builder.class_structure ~self:(Builder.ppat_any ~loc) ~fields:methods))
  in
  let entries_thunk = [%expr fun () -> [%e Builder.elist ~loc entry_exprs]] in
  let register_call =
    Builder.pexp_apply ~loc
      (js_obj_internal_expression ~loc register_name)
      [ (Nolabel, Builder.evar ~loc object_name); (Nolabel, entries_thunk) ]
  in
  let register_call =
    match as_type with None -> register_call | Some core_type -> Builder.pexp_constraint ~loc register_call core_type
  in
  List.fold_right
    (fun binding acc -> Builder.pexp_let ~loc Nonrecursive [ binding ] acc)
    (field_bindings @ [ object_binding ]) register_call

let rec get_return_core_type = function
  | { ptyp_desc = Ptyp_arrow (_, _, rest); _ } -> get_return_core_type rest
  | core_type -> core_type

let get_function_name pattern =
  let rec go pattern =
    match pattern with
    | Ppat_var { txt = name; _ } -> Some name
    | Ppat_constraint (pattern, _) -> go pattern.ppat_desc
    | _ -> None
  in
  go pattern

let get_label = function Ptyp_constr ({ txt = Lident label; _ }, _) -> Some label | _ -> None

(* Extract the `t` from [@mel.send.pipe: t] *)
let get_send_pipe pval_attributes =
  if is_send_pipe pval_attributes then
    let first_attribute = List.hd pval_attributes in
    match first_attribute.attr_payload with PTyp core_type -> Some core_type | _ -> None
  else None

let has_ptyp_attribute ptyp_attributes attribute =
  List.exists (fun { attr_name = { txt = attr } } -> attr = attribute) ptyp_attributes

let is_mel_as core_type =
  match core_type with
  | { ptyp_desc = Ptyp_any; ptyp_attributes; _ } -> has_ptyp_attribute ptyp_attributes "mel.as"
  | _ -> false

let extract_args_labels_types acc pval_type =
  let rec go acc = function
    (* In case of being mel.as, ignore those *)
    | { ptyp_desc = Ptyp_arrow (_label, t1, _t2); _ } when is_mel_as t1 -> acc
    | { ptyp_desc = Ptyp_arrow (_label, _t1, t2); _ } when is_mel_as t2 -> acc
    | { ptyp_desc = Ptyp_arrow (_label, t1, t2); _ } when is_mel_as t1 && is_mel_as t2 -> acc
    | { ptyp_desc = Ptyp_arrow (label, t1, t2); _ } ->
        let pattern = Builder.ppat_var ~loc:t1.ptyp_loc { loc = t1.ptyp_loc; txt = "_" } in
        go ((label, pattern, t1) :: acc) t2
    | _ -> acc
  in
  go acc pval_type

(* Insert send_pipe_core_type as a last argument of the function, but not the return type *)
let construct_pval_with_send_pipe send_pipe_core_type pval_type =
  let rec insert_core_type_in_arrow core_type =
    match core_type with
    (* Handle only ptyp and constr.
       Missing `| Ptyp_any | Ptyp_var | Ptyp_arrow | Ptyp_tuple | Ptyp_constr
                | Ptyp_object | Ptyp_class | Ptyp_alias | Ptyp_variant
                | Ptyp_poly | Ptyp_package | Ptyp_extension`
       The aren't used in most bindings.
    *)
    | { ptyp_desc = Ptyp_arrow (label, t1, t2); _ } -> (
        match (t1.ptyp_desc, t2.ptyp_desc) with
        (* `constr -> arrow (constr -> constr)` gets transformed into
           `constr -> constr -> t -> constr` *)
        | Ptyp_constr _, Ptyp_arrow (_inner_label, _p1, _p2) ->
            Builder.ptyp_arrow ~loc:t1.ptyp_loc label t1 (insert_core_type_in_arrow t2)
        (* `constr -> constr` gets transformed into `constr -> t -> constr` *)
        (* `arrow (constr -> constr) -> constr` gets transformed into,
            `arrow (constr -> constr) -> t -> constr` *)
        | _, _ ->
            Builder.ptyp_arrow ~loc:t2.ptyp_loc label t1
              (Builder.ptyp_arrow ~loc:t2.ptyp_loc Nolabel send_pipe_core_type t2))
    (* In case of being a single ptyp_* turn into ptyp_* -> t *)
    | { ptyp_desc = Ptyp_constr ({ txt = _; loc }, _); _ } | { ptyp_desc = Ptyp_var _; ptyp_loc = loc; _ } ->
        Builder.ptyp_arrow ~loc Nolabel core_type send_pipe_core_type
    (* Here we ignore the Ptyp_any *)
    | _ -> core_type
  in
  insert_core_type_in_arrow pval_type

let inject_send_pipe_as_last_argument pipe_type args_labels =
  match pipe_type with None -> args_labels | Some pipe_core_type -> pipe_core_type :: args_labels

let is_mel_raw expr = match expr with Pexp_extension ({ txt = "mel.raw"; _ }, _) -> true | _ -> false

let capture_payload expr =
  match expr with
  | PStr [ { pstr_desc = Pstr_eval ({ pexp_desc = Pexp_constant (Pconst_string (payload, _, _)); _ }, _); _ } ] ->
      payload
  | _ -> "..."

let get_payload_from_mel_raw expr =
  let rec go expr =
    match expr with
    | Pexp_extension ({ txt = "mel.raw"; _ }, pstr) -> capture_payload pstr
    | Pexp_constraint (expr, _) -> go expr.pexp_desc
    | Pexp_function (_, _, Pfunction_body expr) -> go expr.pexp_desc
    | _ -> "..."
  in
  go expr

let expression_has_mel_raw expr =
  let rec go expr =
    match expr with
    | Pexp_extension ({ txt = "mel.raw"; _ }, _) as pexp_desc -> is_mel_raw pexp_desc
    | Pexp_constraint (expr, _) -> is_mel_raw expr.pexp_desc
    | Pexp_function (_, _, Pfunction_body expr) -> go expr.pexp_desc
    | _ -> false
  in
  go expr

let raise_failure ~loc name =
  [%expr
    let () =
      Printf.printf
        {|
There is a Melange's external (for example: [@mel.get]) call from native code.

Melange externals are bindings to JavaScript code, which can't run on the server and should be wrapped with browser_only ppx or only run it only on the client side. If there's any issue, try wrapping the expression with a try/catch as a workaround.
|}
    in
    raise (Runtime.fail_impossible_action_in_ssr [%e Builder.pexp_constant ~loc (Pconst_string (name, loc, None))])]

let validate_mel_obj_primitive ~loc pval_prim =
  match pval_prim with
  | [] -> ()
  | prims when List.for_all (String.equal "") prims -> ()
  | _ ->
      Location.raise_errorf ~loc
        "[server-reason-react.melange_ppx] [@@mel.obj] requires its external payload to be the empty string"

let transform_external_obj ~loc pval_name pval_type =
  let function_core_type = Builder.ppat_var ~loc:pval_name.loc { loc = pval_name.loc; txt = pval_name.txt } in
  let pat =
    Builder.ppat_constraint ~loc:pval_type.ptyp_loc function_core_type
      (Builder.ptyp_poly ~loc:pval_type.ptyp_loc [] pval_type)
  in
  let rec collect_arguments function_args fields = function
    | { ptyp_desc = Ptyp_arrow (label, core_type, rest); _ } -> (
        if is_mel_as core_type then
          Location.raise_errorf ~loc:core_type.ptyp_loc
            "[server-reason-react.melange_ppx] [@mel.as] is not supported in native [@@mel.obj] externals yet";
        match label with
        | Nolabel ->
            let pattern =
              match core_type.ptyp_desc with
              | Ptyp_constr ({ txt = Lident "unit"; _ }, []) -> Builder.ppat_any ~loc:core_type.ptyp_loc
              | _ ->
                  Location.raise_errorf ~loc:core_type.ptyp_loc
                    "[server-reason-react.melange_ppx] [@@mel.obj] externals in native only support labelled \
                     arguments, optionally labelled arguments, and a final unit argument"
            in
            collect_arguments ((label, pattern, core_type.ptyp_loc) :: function_args) fields rest
        | Labelled name | Optional name ->
            let ident = { loc = core_type.ptyp_loc; txt = name } in
            let pattern = Builder.ppat_var ~loc:core_type.ptyp_loc ident in
            let value = Builder.pexp_ident ~loc:core_type.ptyp_loc { loc = core_type.ptyp_loc; txt = Lident name } in
            let present_expr =
              match label with
              | Labelled _ -> Builder.ebool ~loc:core_type.ptyp_loc true
              | Optional _ -> option_is_some_expr ~loc:core_type.ptyp_loc value
              | Nolabel -> assert false
            in
            let field =
              { method_name = name; js_name = translate_mel_obj_label name; present_expr; value_expr = value }
            in
            collect_arguments ((label, pattern, core_type.ptyp_loc) :: function_args) (field :: fields) rest)
    | _ -> (List.rev function_args, List.rev fields)
  in
  let function_args, fields = collect_arguments [] [] pval_type in
  let object_expression =
    build_registered_js_object_expression ~loc ~register_name:"register_deferred_abstract"
      ~as_type:(get_return_core_type pval_type) fields
  in
  let function_expression =
    List.fold_right
      (fun (label, arg_pat, arg_loc) acc -> Builder.pexp_fun ~loc:arg_loc label None arg_pat acc)
      function_args object_expression
  in
  let vb = Builder.value_binding ~loc ~pat ~expr:function_expression in
  Ast_helper.Str.value Nonrecursive [ vb ]

let mel_raw_found_in_native_message ~loc payload =
  let msg =
    Printf.sprintf
      "[server-reason-react.melange_ppx] There's a [%%mel.raw \"%s\"] expression in native, which should only happen \
       in JavaScript. You need to conditionally run it via let%%browser_only or switch%%platform. More info at \
       https://ml-in-barcelona.github.io/server-reason-react/server-reason-react/browser_ppx.html"
      payload
  in
  Builder.pexp_constant ~loc (Pconst_string (msg, loc, None))

let mel_module_found_in_native_message ~loc =
  let msg =
    Printf.sprintf
      "[server-reason-react.melange_ppx] There's an external with [%%mel.module \"...\"] in native, which should only \
       happen in JavaScript. You need to conditionally discard it from the native build, either by moving the external \
       in a module only available in native, or annotating it with [@platform js]. More info at \
       https://ml-in-barcelona.github.io/server-reason-react/server-reason-react/browser_ppx.html"
  in
  Builder.pexp_constant ~loc (Pconst_string (msg, loc, None))

let external_found_in_native_message ~loc =
  let msg =
    Printf.sprintf
      "[server-reason-react.melange_ppx] There's an external in native, which should only happen in JavaScript. You \
       need to conditionally discard it from the native build, either by moving the external in a module only \
       available in native, or annotating it with [@platform js]. More info at \
       https://ml-in-barcelona.github.io/server-reason-react/server-reason-react/browser_ppx.html"
  in
  Builder.pexp_constant ~loc (Pconst_string (msg, loc, None))

let get_function_arity pattern =
  let rec go arity params body =
    match params with
    | _ :: rest -> go (arity + 1) rest body
    | [] -> (
        match body with
        | Pfunction_body { pexp_desc = Pexp_function (more_params, _, inner_body); _ } ->
            go arity more_params inner_body
        | _ -> arity)
  in
  match pattern with Pexp_function (params, _, body) -> go 0 params body | _ -> 0

let transform_external_arrow ~loc pval_name pval_attributes pval_type =
  let pipe_type =
    match get_send_pipe pval_attributes with
    | Some core_type ->
        let pattern = Builder.ppat_var ~loc:core_type.ptyp_loc { loc = core_type.ptyp_loc; txt = "_" } in
        Some (Nolabel, pattern, core_type)
    | None -> None
  in
  let args_labels_types = extract_args_labels_types [] pval_type in
  let function_core_type = Builder.ppat_var ~loc:pval_name.loc { loc = pval_name.loc; txt = pval_name.txt } in
  let pval_type_piped =
    match pipe_type with
    | None -> pval_type
    | Some (_, _, pipe_type) -> construct_pval_with_send_pipe pipe_type pval_type
  in
  let pat =
    Builder.ppat_constraint ~loc:pval_type.ptyp_loc function_core_type
      (Builder.ptyp_poly ~loc:pval_type.ptyp_loc [] pval_type_piped)
  in
  let arg_labels = inject_send_pipe_as_last_argument pipe_type args_labels_types in
  let function_expression =
    List.fold_left
      (fun acc (label, arg_pat, arg_type) -> Builder.pexp_fun ~loc:arg_type.ptyp_loc label None arg_pat acc)
      (raise_failure ~loc:pval_type.ptyp_loc pval_name.txt)
      arg_labels
  in
  let vb = Builder.value_binding ~loc ~pat ~expr:function_expression in
  Ast_helper.Str.value Nonrecursive [ vb ]

let ptyp_humanize = function
  | Ptyp_tuple _ -> "Tuples"
  | Ptyp_object _ -> "Objects"
  | Ptyp_class _ -> "Classes"
  | Ptyp_variant _ -> "Variants"
  | Ptyp_extension _ -> "Extensions"
  | Ptyp_alias _ -> "Alias"
  | Ptyp_poly _ -> "Polyvariants"
  | Ptyp_package _ -> "Packages"
  | Ptyp_any -> "Any"
  | Ptyp_var _ -> "Var"
  | Ptyp_arrow _ -> "Arrow"
  | Ptyp_constr _ -> "Constr"
  | Ptyp_open _ -> "Open"

let transform_external ~module_path pval_name pval_attributes pval_loc pval_type pval_prim =
  let loc = pval_loc in
  match pval_type.ptyp_desc with
  | Ptyp_arrow _ ->
      if has_attribute pval_attributes "mel.obj" then (
        validate_mel_obj_primitive ~loc pval_prim;
        transform_external_obj ~loc pval_name pval_type)
      else transform_external_arrow ~loc pval_name pval_attributes pval_type
  | Ptyp_var _ | Ptyp_any | Ptyp_constr _ ->
      (* When mel.send.pipe is used, it's treated as a funcion *)
      if Option.is_some (get_send_pipe pval_attributes) then
        transform_external_arrow ~loc pval_name pval_attributes pval_type
      else if Mel_module.has_attr pval_attributes then
        match Mel_module.asset_payload pval_attributes with
        | None ->
            (* If it doesn't have asset payload, we error out as it must be some .js module or package being imported *)
            [%stri [%%ocaml.error [%e mel_module_found_in_native_message ~loc]]]
        | Some str ->
            (* If it has asset payload (file with extension), calculate hash and replace external *)
            let name = Builder.pvar ~loc:pval_name.loc pval_name.txt in
            let path =
              let asset_path = Filename.(concat (dirname module_path) str) in
              let s = In_channel.with_open_bin asset_path In_channel.input_all in
              let filename_fn =
                match !Mel_module.bundler with
                | Webpack -> Mel_module.Webpack.filename
                | Esbuild -> Mel_module.Esbuild.filename
              in
              let prefix = !Mel_module.prefix in
              Builder.estring ~loc Filename.(concat prefix (filename_fn ~base:(Filename.basename str) s))
            in
            [%stri let [%p name] = [%e path]]
      else [%stri [%%ocaml.error [%e external_found_in_native_message ~loc]]]
  | _ ->
      [%stri
        [%%ocaml.error
        "[server-reason-react.melange_ppx] %s are not supported in native externals the same way as melange.ppx \
         support them."
        (ptyp_humanize pval_type.ptyp_desc)]]

let validate_record_labels ~loc record =
  List.fold_left
    (fun acc (longident, expression) ->
      match acc with
      | Error _ as error -> error
      | Ok acc -> (
          match longident.txt with
          | Lident label ->
              Ok
                ({
                   method_name = label;
                   js_name = translate_mel_obj_label label;
                   present_expr = Builder.ebool ~loc:expression.pexp_loc true;
                   value_expr = expression;
                 }
                :: acc)
          | Ldot _ | Lapply _ ->
              Error
                (Location.error_extensionf ~loc
                   "[server-reason-react.melange_ppx] Js.t objects only support labels as keys")))
    (Ok []) record
  |> Result.map List.rev

class raise_exception_mapper (module_path : string) =
  object (_self)
    inherit Ast_traverse.map as super

    method! expression expr =
      let expr = super#expression expr in
      match expr.pexp_desc with
      | Pexp_extension
          ( { txt = "mel.obj"; _ },
            PStr [ { pstr_desc = Pstr_eval ({ pexp_desc = Pexp_record (record, None); pexp_loc }, _); _ } ] ) -> (
          match validate_record_labels ~loc:pexp_loc record with
          | Ok record -> build_registered_js_object_expression ~loc:pexp_loc record
          | Error extension -> Builder.pexp_extension ~loc:pexp_loc extension)
      | Pexp_extension ({ txt = "mel.obj"; loc }, _) ->
          Builder.pexp_extension ~loc
            (Location.error_extensionf ~loc:expr.pexp_loc
               "[server-reason-react.melange_ppx] Js.t objects requires a record literal")
      | Pexp_constant (Pconst_string (s, loc, Some "j")) -> String_interpolation.transform ~loc expr s
      | _ -> expr

    method! structure_item item =
      match item.pstr_desc with
      (* [%%mel.raw ...] *)
      | Pstr_extension (({ txt = "mel.raw"; _ }, pstr), _) ->
          let loc = item.pstr_loc in
          let payload = capture_payload pstr in
          [%stri [%%ocaml.error [%e mel_raw_found_in_native_message ~loc payload]]]
      (* let a _ = [%mel.raw ...] *)
      | Pstr_value
          ( Nonrecursive,
            [
              {
                pvb_expr = { pexp_desc = Pexp_function (_ :: _, _, Pfunction_body expression); _ };
                pvb_pat = { ppat_desc = Ppat_var { txt = _function_name; _ } };
                pvb_attributes = _;
                pvb_loc;
              };
            ] )
        when expression_has_mel_raw expression.pexp_desc ->
          let loc = item.pstr_loc in
          let payload = get_payload_from_mel_raw expression.pexp_desc in
          [%stri [%error [%e mel_raw_found_in_native_message ~loc:pvb_loc payload]]]
      (* let a = [%mel.raw ...] *)
      | Pstr_value
          ( Nonrecursive,
            [
              {
                pvb_expr = expression;
                pvb_pat = { ppat_desc = Ppat_var { txt = _function_name; _ } };
                pvb_attributes = _;
                pvb_loc;
              };
            ] )
        when expression_has_mel_raw expression.pexp_desc ->
          let loc = item.pstr_loc in
          let payload = get_payload_from_mel_raw expression.pexp_desc in
          [%stri [%error [%e mel_raw_found_in_native_message ~loc:pvb_loc payload]]]
      (* let a: t = [%mel.raw ...] *)
      | Pstr_value
          (Nonrecursive, [ { pvb_expr = expression; pvb_pat = { ppat_desc = _ }; pvb_attributes = _; pvb_loc } ])
        when expression_has_mel_raw expression.pexp_desc ->
          let loc = item.pstr_loc in
          let payload = get_payload_from_mel_raw expression.pexp_desc in
          [%stri [%error [%e mel_raw_found_in_native_message ~loc:pvb_loc payload]]]
      (* %mel. *)
      (* external foo: t = "{{JavaScript}}" *)
      | Pstr_primitive { pval_name; pval_attributes; pval_loc; pval_type; pval_prim } ->
          (* Detects [@browser_only] or [@platform js] attributes. When present on an external, we pass it through unchanged so browser_ppx can filter it out in native mode. *)
          if has_browser_ppx_attribute pval_attributes then item
          else transform_external ~module_path pval_name pval_attributes pval_loc pval_type pval_prim
      | _ -> super#structure_item item
  end

let structure_mapper ctxt s =
  let module_path = Code_path.file_path (Expansion_context.Base.code_path ctxt) in
  (new raise_exception_mapper module_path)#structure s

module Debug = struct
  let rule =
    let extractor = Ast_pattern.(__') in
    let handler ~ctxt:_ { loc } = [%expr ()] in
    Context_free.Rule.extension (Extension.V3.declare "debug" Extension.Context.expression extractor handler)
end

let () =
  Driver.add_arg "-bundler"
    (String
       (fun str ->
         match str with
         | "webpack" -> Mel_module.bundler := Webpack
         | "esbuild" -> Mel_module.bundler := Esbuild
         | _ ->
             failwith
               (Printf.sprintf
                  {|Unknown value %S passed as -bundler flag in melange.ppx, valid values: "webpack", "esbuild"|} str)))
    ~doc:"generate paths to assets in mel.module using the file name scheme of the bundler of choice";
  Driver.add_arg "-prefix"
    (String (fun str -> Mel_module.prefix := str))
    ~doc:"the paths to the generated assets will include the given prefix before the filename (default: \"/\")";
  Driver.V2.register_transformation ~impl:structure_mapper
    ~rules:[ Pipe_first.rule; Regex.rule; Double_hash.rule; Debug.rule; Private.rule ]
    "melange-native-ppx"