Source file static_analysis.ml

open Ppxlib

type static_part =
  | Static_str of string
  | Dynamic_string of expression
  | Dynamic_int of expression
  | Dynamic_float of expression
  | Dynamic_stringf of expression * (arg_label * expression) list
  | Dynamic_element of expression

let rec coalesce_static_parts = function
  | Static_str a :: Static_str b :: rest ->
      coalesce_static_parts (Static_str (a ^ b) :: rest)
  | x :: rest ->
      x :: coalesce_static_parts rest
  | [] ->
      []

let escape_html s =
  let len = String.length s in
  let rec find_escape i =
    if i = len then
      None
    else
      match s.[i] with
      | '&' | '<' | '>' | '\'' | '"' ->
          Some i
      | _ ->
          find_escape (i + 1)
  in
  match find_escape 0 with
  | None ->
      s
  | Some first ->
      let buf = Buffer.create (len * 2) in
      if first > 0 then Buffer.add_substring buf s 0 first;
      for i = first to len - 1 do
        match s.[i] with
        | '&' ->
            Buffer.add_string buf "&amp;"
        | '<' ->
            Buffer.add_string buf "&lt;"
        | '>' ->
            Buffer.add_string buf "&gt;"
        | '\'' ->
            Buffer.add_string buf "&apos;"
        | '"' ->
            Buffer.add_string buf "&quot;"
        | c ->
            Buffer.add_char buf c
      done;
      Buffer.contents buf

(* Duplicated from JSX.Html.is_self_closing_tag - keep in sync *)
let is_self_closing_tag = function
  | "area"
  | "base"
  | "br"
  | "col"
  | "embed"
  | "hr"
  | "img"
  | "input"
  | "link"
  | "meta"
  | "param"
  | "source"
  | "track"
  | "wbr"
  | "menuitem" ->
      true
  | _ ->
      false

let rec extract_literal_string expr =
  match expr.pexp_desc with
  | Pexp_constant (Pconst_string (s, _, _)) ->
      Some s
  | Pexp_constraint (inner, _) ->
      extract_literal_string inner
  | _ ->
      None

let rec extract_literal_int expr =
  match expr.pexp_desc with
  | Pexp_constant (Pconst_integer (s, _)) ->
      Some (int_of_string s)
  | Pexp_constraint (inner, _) ->
      extract_literal_int inner
  | _ ->
      None

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

let extract_jsx_string_arg expr =
  match expr.pexp_desc with
  | Pexp_apply
      ( {
          pexp_desc =
            Pexp_ident { txt = Ldot (Lident "JSX", ("text" | "string")); _ };
          _;
        },
        [ (Nolabel, arg) ]
      )
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Lident ("text" | "string"); _ }; _ },
        [ (Nolabel, arg) ]
      ) ->
      Some arg
  | _ ->
      None

let extract_jsx_int_arg expr =
  match expr.pexp_desc with
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Ldot (Lident "JSX", "int"); _ }; _ },
        [ (Nolabel, arg) ]
      )
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Lident "int"; _ }; _ },
        [ (Nolabel, arg) ]
      ) ->
      Some arg
  | _ ->
      None

let extract_jsx_float_arg expr =
  match expr.pexp_desc with
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Ldot (Lident "JSX", "float"); _ }; _ },
        [ (Nolabel, arg) ]
      )
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Lident "float"; _ }; _ },
        [ (Nolabel, arg) ]
      ) ->
      Some arg
  | _ ->
      None

let extract_jsx_stringf_call expr =
  match expr.pexp_desc with
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Ldot (Lident "JSX", "stringf"); _ }; _ },
        (Nolabel, fmt) :: args
      )
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Lident "stringf"; _ }; _ },
        (Nolabel, fmt) :: args
      ) ->
      Some (fmt, args)
  | _ ->
      None

let rec extract_literal_char expr =
  match expr.pexp_desc with
  | Pexp_constant (Pconst_char c) ->
      Some c
  | Pexp_constraint (inner, _) ->
      extract_literal_char inner
  | _ ->
      None

type literal_format_arg =
  | Literal_format_string of string
  | Literal_format_char of char
  | Literal_format_int of int

let extract_literal_format_arg = function
  | Nolabel, expr -> (
      match extract_literal_string expr with
      | Some s ->
          Some (Literal_format_string s)
      | None -> (
          match extract_literal_char expr with
          | Some c ->
              Some (Literal_format_char c)
          | None ->
              extract_literal_int expr
              |> Option.map (fun i -> Literal_format_int i)
        )
    )
  | _ ->
      None

let render_literal_stringf fmt args =
  let len = String.length fmt in
  let buf = Buffer.create len in
  let rec loop i remaining_args =
    if i >= len then
      match remaining_args with [] -> Some (Buffer.contents buf) | _ -> None
    else if String.unsafe_get fmt i = '%' then
      if i + 1 >= len then
        None
      else
        match String.unsafe_get fmt (i + 1) with
        | '%' ->
            Buffer.add_char buf '%';
            loop (i + 2) remaining_args
        | 's' -> (
            match remaining_args with
            | arg :: rest -> (
                match extract_literal_format_arg arg with
                | Some (Literal_format_string s) ->
                    Buffer.add_string buf s;
                    loop (i + 2) rest
                | _ ->
                    None
              )
            | [] ->
                None
          )
        | 'c' -> (
            match remaining_args with
            | arg :: rest -> (
                match extract_literal_format_arg arg with
                | Some (Literal_format_char c) ->
                    Buffer.add_char buf c;
                    loop (i + 2) rest
                | _ ->
                    None
              )
            | [] ->
                None
          )
        | 'd' | 'i' -> (
            match remaining_args with
            | arg :: rest -> (
                match extract_literal_format_arg arg with
                | Some (Literal_format_int value) ->
                    Buffer.add_string buf (string_of_int value);
                    loop (i + 2) rest
                | _ ->
                    None
              )
            | [] ->
                None
          )
        | _ ->
            None
    else
      let c = String.unsafe_get fmt i in
      Buffer.add_char buf c;
      loop (i + 1) remaining_args
  in
  loop 0 args

let extract_jsx_stringf_literal expr =
  match extract_jsx_stringf_call expr with
  | Some (fmt_expr, args) -> (
      match extract_literal_string fmt_expr with
      | Some fmt ->
          render_literal_stringf fmt args |> Option.map escape_html
      | None ->
          None
    )
  | None ->
      None

let extract_jsx_text_literal expr =
  match extract_jsx_string_arg expr with
  | Some arg ->
      extract_literal_string arg
  | None ->
      None

type static_attr_value =
  | Static_string of string
  | Static_int of int
  | Static_bool of bool

let extract_static_attr_value expr =
  match extract_literal_string expr with
  | Some s ->
      Some (Static_string s)
  | None -> (
      match extract_literal_int expr with
      | Some i ->
          Some (Static_int i)
      | None -> (
          match extract_literal_bool expr with
          | Some b ->
              Some (Static_bool b)
          | None ->
              None
        )
    )

let render_attr_value = function
  | Static_string s ->
      escape_html s
  | Static_int i ->
      string_of_int i
  | Static_bool true ->
      "true"
  | Static_bool false ->
      "false"

type attr_render_info = {
  html_name : string;
  is_boolean : bool;
  kind : Html_attributes.kind;
}

type parsed_attr =
  | Static_attr of attr_render_info * static_attr_value
  | Optional_attr of attr_render_info * expression
  | Dynamic_attr of attr_render_info * expression

type attr_validation_result = Valid_attr of attr_render_info | Invalid_attr

let validate_attr_for_static ~tag_name jsx_name =
  match Html.findByName tag_name jsx_name with
  | Error _ ->
      Invalid_attr
  | Ok prop ->
      let html_name = Html.getName prop in
      let kind =
        match prop with
        | Html_attributes.Attribute { type_; _ }
        | Html_attributes.Rich_attribute { type_; _ } ->
            type_
        | Html_attributes.Event _ ->
            Html_attributes.String
      in
      let is_boolean = kind = Html_attributes.Bool in
      Valid_attr { html_name; is_boolean; kind }

let render_static_attr_with_info info value =
  match value with
  | Static_bool false when info.is_boolean ->
      ""
  | Static_bool true when info.is_boolean ->
      " " ^ info.html_name
  | _ ->
      let value_str = render_attr_value value in
      Printf.sprintf " %s=\"%s\"" info.html_name value_str

type attr_analysis_result = Ok of parsed_attr option | Invalid

let analyze_attribute ~tag_name (label, expr) : attr_analysis_result =
  match label with
  | Nolabel ->
      Ok None (* Children, handled separately *)
  | Optional name -> (
      match validate_attr_for_static ~tag_name name with
      | Invalid_attr ->
          Invalid
      | Valid_attr info ->
          Ok (Some (Optional_attr (info, expr)))
    )
  | Labelled name -> (
      match validate_attr_for_static ~tag_name name with
      | Invalid_attr ->
          Invalid
      | Valid_attr info -> (
          let static_value =
            match info.kind with
            | Html_attributes.Polyvariant options -> (
                match expr.pexp_desc with
                | Pexp_variant (constructor, None) -> (
                    match
                      List.find_opt
                        (fun (opt : Html_attributes.polyvariant) ->
                          opt.type_ = constructor
                        )
                        options
                    with
                    | Some opt ->
                        Some (Static_string opt.jsxName)
                    | None ->
                        None
                  )
                | _ ->
                    None
              )
            | _ ->
                extract_static_attr_value expr
          in
          match static_value with
          | Some value ->
              Ok (Some (Static_attr (info, value)))
          | None ->
              Ok (Some (Dynamic_attr (info, expr)))
        )
    )

type attrs_analysis =
  | All_static of string
  | Has_optional of (attr_render_info * expression) list * string
  | Has_dynamic_attrs of {
      static_attrs : string;
      dynamic_attrs : (attr_render_info * expression) list;
    }
  | Validation_failed

let analyze_attributes ~tag_name attrs =
  let rec loop static_buf optionals dynamic_attrs = function
    | [] ->
        if dynamic_attrs <> [] then
          Has_dynamic_attrs
            {
              static_attrs = Buffer.contents static_buf;
              dynamic_attrs = List.rev dynamic_attrs;
            }
        else if optionals = [] then
          All_static (Buffer.contents static_buf)
        else
          Has_optional (List.rev optionals, Buffer.contents static_buf)
    | attr :: rest -> (
        match analyze_attribute ~tag_name attr with
        | Invalid ->
            Validation_failed
        | Ok None ->
            loop static_buf optionals dynamic_attrs rest
        | Ok (Some (Static_attr (info, value))) ->
            Buffer.add_string static_buf
              (render_static_attr_with_info info value);
            loop static_buf optionals dynamic_attrs rest
        | Ok (Some (Optional_attr (info, expr))) ->
            loop static_buf ((info, expr) :: optionals) dynamic_attrs rest
        | Ok (Some (Dynamic_attr (info, expr))) ->
            loop static_buf optionals ((info, expr) :: dynamic_attrs) rest
      )
  in
  loop (Buffer.create 64) [] [] attrs

type children_analysis =
  | No_children
  | All_static_children of string
  | All_string_dynamic of static_part list
  | Mixed_children of static_part list

let extract_jsx_unsafe_literal expr =
  match expr.pexp_desc with
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Ldot (Lident "JSX", "unsafe"); _ }; _ },
        [ (Nolabel, arg) ]
      ) ->
      extract_literal_string arg
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Lident "unsafe"; _ }; _ },
        [ (Nolabel, arg) ]
      ) ->
      extract_literal_string arg
  | _ ->
      None

let extract_jsx_int_literal expr =
  match extract_jsx_int_arg expr with
  | Some arg ->
      extract_literal_int arg
  | None ->
      None

let extract_jsx_float_literal expr =
  match expr.pexp_desc with
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Ldot (Lident "JSX", "float"); _ }; _ },
        [ (Nolabel, arg) ]
      )
  | Pexp_apply
      ( { pexp_desc = Pexp_ident { txt = Lident "float"; _ }; _ },
        [ (Nolabel, arg) ]
      ) -> (
      match arg.pexp_desc with
      | Pexp_constant (Pconst_float (s, _)) ->
          Some (float_of_string s)
      | Pexp_constraint (inner, _) -> (
          match inner.pexp_desc with
          | Pexp_constant (Pconst_float (s, _)) ->
              Some (float_of_string s)
          | _ ->
              None
        )
      | _ ->
          None
    )
  | _ ->
      None

let analyze_child (expr : expression) : static_part =
  List.find_map
    (fun fn -> fn ())
    [
      (fun () ->
        extract_jsx_unsafe_literal expr |> Option.map (fun s -> Static_str s)
      );
      (fun () ->
        extract_jsx_text_literal expr
        |> Option.map (fun s -> Static_str (escape_html s))
      );
      (fun () ->
        extract_literal_string expr
        |> Option.map (fun s -> Static_str (escape_html s))
      );
      (fun () ->
        extract_jsx_int_literal expr
        |> Option.map (fun i -> Static_str (string_of_int i))
      );
      (fun () ->
        extract_jsx_float_literal expr
        |> Option.map (fun f -> Static_str (Float.to_string f))
      );
      (fun () ->
        extract_jsx_stringf_literal expr |> Option.map (fun s -> Static_str s)
      );
      (fun () ->
        extract_jsx_stringf_call expr
        |> Option.map (fun (fmt, args) -> Dynamic_stringf (fmt, args))
      );
      (fun () ->
        extract_jsx_string_arg expr |> Option.map (fun e -> Dynamic_string e)
      );
      (fun () -> extract_jsx_int_arg expr |> Option.map (fun e -> Dynamic_int e));
      (fun () ->
        extract_jsx_float_arg expr |> Option.map (fun e -> Dynamic_float e)
      );
    ]
  |> Option.value ~default:(Dynamic_element expr)

let analyze_children children =
  match children with
  | None ->
      No_children
  | Some [] ->
      No_children
  | Some children ->
      let parts = List.map analyze_child children in
      let all_static =
        List.for_all (function Static_str _ -> true | _ -> false) parts
      in
      let has_element_dynamic =
        List.exists (function Dynamic_element _ -> true | _ -> false) parts
      in
      if all_static then (
        let buf = Buffer.create 128 in
        List.iter
          (function Static_str s -> Buffer.add_string buf s | _ -> ())
          parts;
        All_static_children (Buffer.contents buf)
      ) else if not has_element_dynamic then
        All_string_dynamic (coalesce_static_parts parts)
      else
        Mixed_children (coalesce_static_parts parts)

type element_analysis =
  | Fully_static of string
  | Needs_string_concat of static_part list
  | Needs_buffer of {
      parts : static_part list;
      static_size : int;
      dynamic_count : int;
    }
  | Has_optional_attrs of {
      tag_name : string;
      static_attrs : string;
      optional_attrs : (attr_render_info * expression) list;
      children_parts : static_part list;
      is_self_closing : bool;
    }
  | Dynamic_attrs_children of {
      tag_name : string;
      static_attrs : string;
      dynamic_attrs : (attr_render_info * expression) list;
      children_parts : static_part list;
      is_self_closing : bool;
    }
  | Cannot_optimize

let analyze_element ~tag_name ~attrs ~children =
  let attrs_result = analyze_attributes ~tag_name attrs in
  let children_result = analyze_children children in

  match (attrs_result, children_result) with
  | Validation_failed, _ ->
      Cannot_optimize
  | ( Has_dynamic_attrs { static_attrs; dynamic_attrs },
      All_static_children children_html ) ->
      Dynamic_attrs_children
        {
          tag_name;
          static_attrs;
          dynamic_attrs;
          children_parts = [ Static_str children_html ];
          is_self_closing = is_self_closing_tag tag_name;
        }
  | Has_dynamic_attrs { static_attrs; dynamic_attrs }, No_children ->
      Dynamic_attrs_children
        {
          tag_name;
          static_attrs;
          dynamic_attrs;
          children_parts = [];
          is_self_closing = is_self_closing_tag tag_name;
        }
  | Has_dynamic_attrs { static_attrs; dynamic_attrs }, All_string_dynamic parts
    ->
      Dynamic_attrs_children
        {
          tag_name;
          static_attrs;
          dynamic_attrs;
          children_parts = parts;
          is_self_closing = false;
        }
  | Has_dynamic_attrs { static_attrs; dynamic_attrs }, Mixed_children parts ->
      Dynamic_attrs_children
        {
          tag_name;
          static_attrs;
          dynamic_attrs;
          children_parts = parts;
          is_self_closing = false;
        }
  | All_static attrs_html, No_children when is_self_closing_tag tag_name ->
      let html = Printf.sprintf "<%s%s />" tag_name attrs_html in
      Fully_static html
  | All_static attrs_html, No_children ->
      let html = Printf.sprintf "<%s%s></%s>" tag_name attrs_html tag_name in
      Fully_static html
  | All_static attrs_html, All_static_children children_html ->
      let html =
        Printf.sprintf "<%s%s>%s</%s>" tag_name attrs_html children_html
          tag_name
      in
      Fully_static html
  | All_static attrs_html, All_string_dynamic parts ->
      let open_tag = Printf.sprintf "<%s%s>" tag_name attrs_html in
      let close_tag = Printf.sprintf "</%s>" tag_name in
      let all_parts =
        [ Static_str open_tag ] @ parts @ [ Static_str close_tag ]
      in
      Needs_string_concat (coalesce_static_parts all_parts)
  | All_static attrs_html, Mixed_children parts ->
      let open_tag = Printf.sprintf "<%s%s>" tag_name attrs_html in
      let close_tag = Printf.sprintf "</%s>" tag_name in
      let all_parts =
        coalesce_static_parts
          ([ Static_str open_tag ] @ parts @ [ Static_str close_tag ])
      in
      let static_size =
        List.fold_left
          (fun acc part ->
            match part with Static_str s -> acc + String.length s | _ -> acc
          )
          0 all_parts
      in
      let dynamic_count =
        List.fold_left
          (fun acc part -> match part with Static_str _ -> acc | _ -> acc + 1)
          0 all_parts
      in
      Needs_buffer { parts = all_parts; static_size; dynamic_count }
  | Has_optional (optional_attrs, static_attrs), No_children ->
      Has_optional_attrs
        {
          tag_name;
          static_attrs;
          optional_attrs;
          children_parts = [];
          is_self_closing = is_self_closing_tag tag_name;
        }
  | ( Has_optional (optional_attrs, static_attrs),
      All_static_children children_html ) ->
      Has_optional_attrs
        {
          tag_name;
          static_attrs;
          optional_attrs;
          children_parts = [ Static_str children_html ];
          is_self_closing = false;
        }
  | Has_optional (optional_attrs, static_attrs), All_string_dynamic parts ->
      Has_optional_attrs
        {
          tag_name;
          static_attrs;
          optional_attrs;
          children_parts = parts;
          is_self_closing = false;
        }
  | Has_optional (optional_attrs, static_attrs), Mixed_children parts ->
      Has_optional_attrs
        {
          tag_name;
          static_attrs;
          optional_attrs;
          children_parts = parts;
          is_self_closing = false;
        }

let maybe_add_doctype tag_name html =
  if tag_name = "html" then
    "<!DOCTYPE html>" ^ html
  else
    html