Improve parse performance, and fix quadratic runtime on deeply nested elements

src/common.ml

@@ -137,7 +137,11 @@ let char c =
   else
     format_char c
 
-let is_valid_html_char c = not (is_control_character c || is_non_character c)
+let is_valid_html_char c =
+  (* Fast path for printable ASCII, which is the vast majority of input and is
+     always valid; avoids two function calls per character on the hot path. *)
+  is_in_range 0x0020 0x007E c
+  || not (is_control_character c || is_non_character c)
 
 let is_valid_xml_char c =
   is_in_range 0x0020 0xD7FF c

src/encoding.ml

@@ -10,11 +10,20 @@ let wrap f = fun ?(report = Error.ignore_errors) s -> f report s
 
 let bytes_empty = Bytes.create 0
 
+let chunk_size = 4096
+
 (* Decoders based on the Uutf library. *)
 let uutf_decoder encoding name =
   (fun report bytes ->
     let decoder = Uutf.decoder ~encoding `Manual in
 
+    (* The byte source yields one byte per pull, but feeding Uutf a byte at a
+       time is slow. Instead we buffer up to [chunk_size] bytes and let Uutf
+       decode the whole chunk before refilling. The buffer is reused across
+       chunks; Uutf consumes each chunk fully (copying any straddling sequence
+       into its own scratch) before returning `Await, so overwriting it is safe. *)
+    let chunk = Bytes.create chunk_size in
+
     (fun throw empty k ->
       let rec run () =
         match Uutf.decode decoder with
@@ -24,10 +33,22 @@ let uutf_decoder encoding name =
           let location = Uutf.decoder_line decoder, Uutf.decoder_col decoder in
           report location (`Decoding_error (s, name)) throw (fun () ->
           k u_rep)
-        | `Await ->
+        | `Await -> fill 0
+
+      (* Fill [chunk] with bytes and then handoff to [feed]. *)
+      and fill n =
+        if n >= chunk_size then feed n
+        else
           next bytes throw
-            (fun () -> Uutf.Manual.src decoder bytes_empty 0 0; run ())
-            (fun c -> Uutf.Manual.src decoder (Bytes.make 1 c) 0 1; run ())
+            (fun () -> feed n)
+            (* unsafe_set is safe because of the above chunk_size guard. *)
+            (fun c -> Bytes.unsafe_set chunk n c; fill (n + 1))
+
+      (* Feed [chunk] to Uutf and then handoff to [run]. *)
+      and feed n =
+        if n = 0 then Uutf.Manual.src decoder bytes_empty 0 0
+        else Uutf.Manual.src decoder chunk 0 n;
+        run ()
       in
       run ())
     |> make)

src/html_parser.ml

@@ -1027,6 +1027,23 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
     report l (`Misnested_tag (t.name, context_name, t.Token_tag.attributes)) !throw k in
 
   let open_elements = Stack.create () in
+
+  (* Number of open <p> elements on the stack. [in_button_scope "p"] can only be
+     true when this is positive, so tracking it lets close_current_p_element skip
+     a full open-elements scan on every block-level start tag. Without it,
+     opening N nested block elements (e.g. <div>) is O(N^2), because each open
+     scans the whole growing stack for a p to close. Maintained incrementally on
+     push/pop; recomputed after the (rare) adoption agency, which can also remove
+     p elements. *)
+  let open_p_count = ref 0 in
+  let is_paragraph element =
+    match element.element_name with `HTML, "p" -> true | _ -> false in
+  let count_open_p () =
+    List.fold_left
+      (fun n element -> if is_paragraph element then n + 1 else n)
+      0 !open_elements
+  in
+
   let active_formatting_elements = Active.create () in
   let subtree_buffer = Subtree.create open_elements in
   let text = Text.prepare () in
@@ -1202,6 +1219,7 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
       Element.create ~is_html_integration_point (namespace, name) location
     in
     open_elements := element_entry::!open_elements;
+    if is_paragraph element_entry then incr open_p_count;
 
     if set_form_element_pointer then
       form_element_pointer := Some element_entry;
@@ -1229,6 +1247,7 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
       (fun () ->
         open_elements := more;
         element.is_open <- false;
+        if is_paragraph element then decr open_p_count;
         if element.suppress then mode ()
         else emit' location `End_element mode))
 
@@ -1311,8 +1330,18 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
       | _ -> false) location (fun () ->
     pop location mode)))
 
+  (* The adoption agency can remove p elements from the stack (they are not
+     formatting elements, but can lie between one and its furthest block), so
+     recompute the cached count afterwards rather than trying to track its
+     internal stack surgery. It is rare and already O(stack), so the extra scan
+     is immaterial. *)
+  and run_adoption_agency l name mode =
+    adoption_agency_algorithm l name (fun () ->
+      open_p_count := count_open_p ();
+      mode ())
+
   and close_current_p_element l mode =
-    if Stack.in_button_scope open_elements "p" then
+    if !open_p_count > 0 && Stack.in_button_scope open_elements "p" then
       close_element_with_implied "p" l mode
     else mode ()
 
@@ -1338,6 +1367,16 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
     !ended ()))
 
   and reconstruct_active_formatting_elements mode =
+    match !active_formatting_elements with
+    (* Fast path: there is nothing to reconstruct when the list is empty or its
+       most recent entry is a marker or an already-open element. This holds for
+       essentially every character insertion, which calls this per character, so
+       skip the list walk, tuple allocation, and ref write below. *)
+    | []
+    | Active.Marker::_
+    | Active.Element_ ({is_open = true}, _, _)::_ -> mode ()
+
+    | _ ->
     let rec get_prefix prefix = function
       | [] -> prefix, []
       | Active.Marker::_ as l -> prefix, l
@@ -1843,7 +1882,7 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
         | None -> k ()
         | Some existing ->
           misnested_tag l t "a" (fun () ->
-          adoption_agency_algorithm l "a" (fun () ->
+          run_adoption_agency l "a" (fun () ->
           Stack.remove open_elements existing;
           Active.remove active_formatting_elements existing;
           k ())))
@@ -1866,14 +1905,14 @@ let parse requested_context report (tokens, set_tokenizer_state, set_foreign) =
         if not @@ Stack.in_scope open_elements "nobr" then k ()
         else
           misnested_tag l t "nobr" (fun () ->
-          adoption_agency_algorithm l "nobr" (fun () ->
+          run_adoption_agency l "nobr" (fun () ->
           reconstruct_active_formatting_elements k)))
       (fun () -> push_and_emit ~formatting:true l t mode))
 
     | l, `End {name =
         "a" | "b" | "big" | "code" | "em" | "font" | "i" | "nobr" | "s" |
         "small" | "strike" | "strong" | "tt" | "u" as name} ->
-      adoption_agency_algorithm l name mode
+      run_adoption_agency l name mode
 
     | l, `Start ({name = "applet" | "marquee" | "object"} as t) ->
       frameset_ok := false;

src/stream_io.ml

@@ -3,20 +3,24 @@
 
 open Kstream
 
-let state_fold f initial =
-  let state = ref initial in
-  (fun throw e k ->
-    f !state throw e (fun (c, new_state) ->
-      state := new_state; k c))
-  |> make
-
+(* Track the read position with a mutable index, avoiding a per-byte allocation
+   on this hot path. *)
 let string s =
-  state_fold (fun i _ e k ->
-    if i >= String.length s then e () else k (s.[i], i + 1)) 0
+  let position = ref 0 in
+  (fun _ e k ->
+    let i = !position in
+    if i >= String.length s then e ()
+    else (position := i + 1; k (String.unsafe_get s i)))
+  |> make
 
+(* Same as [string], over a Buffer.t. *)
 let buffer b =
-  state_fold (fun i _ e k ->
-    if i >= Buffer.length b then e () else k (Buffer.nth b i, i + 1)) 0
+  let position = ref 0 in
+  (fun _ e k ->
+    let i = !position in
+    if i >= Buffer.length b then e ()
+    else (position := i + 1; k (Buffer.nth b i)))
+  |> make
 
 (* Optimized away by Flambda. *)
 type result = Count of int | Exn of exn

test/performance/performance_common.ml

@@ -1,22 +1,69 @@
 (* This file is part of Markup.ml, released under the MIT license. See
    LICENSE.md for details, or visit https://github.com/aantron/markup.ml. *)
 
+(* Report the median of [runs] timed runs. The median is more stable than the
+   mean under system noise. [f] is expected to parse an already in-memory
+   string, so file I/O is excluded from the measurement. *)
 let measure runs library source format f =
   let name = Printf.sprintf "%s: %s (%s)" library source format in
 
-  let rec run = function
-    | 0 -> ()
-    | n -> f (); run (n - 1)
-  in
+  (* Warm up by running [f] a few times before measuring. *)
+  for _ = 1 to max 1 (runs / 10) do f () done;
 
-  let start_time = Unix.gettimeofday () in
+  let times = Array.make runs 0. in
+  for i = 0 to runs - 1 do
+    let start_time = Unix.gettimeofday () in
+    f ();
+    times.(i) <- ((Unix.gettimeofday ()) -. start_time) *. 1000000.
+  done;
+  Array.sort compare times;
 
-  run runs;
+  Printf.printf "  %s: %.0f us\n" name times.(runs / 2)
 
-  let duration = (Unix.gettimeofday ()) -. start_time in
-  let average = duration /. (float_of_int runs) *. 1000000. in
+let read_file path =
+  let channel = open_in_bin path in
+  let length = in_channel_length channel in
+  let content = really_input_string channel length in
+  close_in channel;
+  content
 
-  Printf.printf "  %s: %.0f us\n" name average
+let google_page = read_file "test/pages/google"
+let xml_spec = read_file "test/pages/xml_spec"
 
-let google_page = "test/pages/google"
-let xml_spec = "test/pages/xml_spec"
+(* The documents below each stress one part of the HTML parser. They are built
+   here by repeating a small pattern to a demanding size, rather than committed
+   as large files, since the pattern is all that matters. *)
+
+let repeat n piece =
+  let buffer = Buffer.create (n * String.length piece) in
+  for _ = 1 to n do Buffer.add_string buffer piece done;
+  Buffer.contents buffer
+
+let html_document body =
+  "<!doctype html>\n<html><head><title>stress</title></head><body>\n"
+  ^ body ^ "\n</body></html>\n"
+
+(* Multi-byte UTF-8 (CJK), which exercises the decoder's multi-byte path. *)
+let stress_cjk =
+  html_document
+    (repeat 2500
+      "<p>甀曒檃檑糲蘥蠩櫋瀩嗢剆坲姏齸圞趲葠蜄蛖砎粁擙樲橚噅尰崺廘榙榾</p>\n")
+
+(* Formatting elements, some deliberately misnested, which exercise the
+   adoption agency algorithm. *)
+let stress_formatting =
+  html_document
+    (repeat 1500
+      "<p><a href=\"/x\">lorem</a> <em>ipsum</em> <strong>dolor</strong> \
+       <b><i>sit</b> amet</i></p>\n")
+
+(* Named, decimal, and hexadecimal character references. *)
+let stress_entities =
+  html_document
+    (repeat 1500
+      "<p>lorem&amp;ipsum&#8226;dolor&#x2022;sit&copy;amet&mdash;elit&nbsp;n</p>\n")
+
+(* A single deep chain of block elements. *)
+let stress_deep_nesting =
+  let depth = 700 in
+  html_document (repeat depth "<div>" ^ "<p>content</p>" ^ repeat depth "</div>")

test/performance/performance_markup.ml

@@ -6,9 +6,23 @@ open Markup
 
 let (|>) x f = f x
 
+let parse_html_string s = string s |> parse_html |> signals |> drain
+
 let () =
-  measure 100 "markup.ml" google_page "html" (fun () ->
-    file google_page |> fst |> parse_html |> signals |> drain);
+  measure 100 "markup.ml" "google" "html"
+    (fun () -> parse_html_string google_page);
+
+  measure 100 "markup.ml" "xml_spec" "xml"
+    (fun () -> string xml_spec |> parse_xml |> signals |> drain);
+
+  measure 100 "markup.ml" "stress_cjk" "html"
+    (fun () -> parse_html_string stress_cjk);
+
+  measure 100 "markup.ml" "stress_formatting" "html"
+    (fun () -> parse_html_string stress_formatting);
+
+  measure 100 "markup.ml" "stress_entities" "html"
+    (fun () -> parse_html_string stress_entities);
 
-  measure 100 "markup.ml" xml_spec "xml" (fun () ->
-    file xml_spec |> fst |> parse_xml |> signals |> drain)
+  measure 100 "markup.ml" "stress_deep_nesting" "html"
+    (fun () -> parse_html_string stress_deep_nesting)

test/performance/performance_nethtml.ml

@@ -6,16 +6,27 @@ open Nethtml
 
 let (|>) x f = f x
 
-let parse file =
-  file
-  |> open_in
-  |> Lexing.from_channel
+let parse s =
+  s
+  |> Lexing.from_string
   |> parse_document ~dtd:relaxed_html40_dtd
   |> ignore
 
 let () =
-  measure 100 "nethtml" google_page "html" (fun () ->
-    parse google_page);
+  measure 100 "nethtml" "google" "html"
+    (fun () -> parse google_page);
 
-  measure 100 "nethtml" xml_spec "html" (fun () ->
-    parse xml_spec)
+  measure 100 "nethtml" "xml_spec" "html"
+    (fun () -> parse xml_spec);
+
+  measure 100 "nethtml" "stress_cjk" "html"
+    (fun () -> parse stress_cjk);
+
+  measure 100 "nethtml" "stress_formatting" "html"
+    (fun () -> parse stress_formatting);
+
+  measure 100 "nethtml" "stress_entities" "html"
+    (fun () -> parse stress_entities);
+
+  measure 100 "nethtml" "stress_deep_nesting" "html"
+    (fun () -> parse stress_deep_nesting)

test/performance/performance_xmlm.ml

@@ -6,15 +6,14 @@ open Xmlm
 
 let (|>) x f = f x
 
-let parse file =
+let parse s =
   try
-    make_input ~entity:(fun _ -> Some "") (`Channel (open_in file))
+    make_input ~entity:(fun _ -> Some "") (`String (0, s))
     |> input_doc_tree ~el:(fun _ _ -> ()) ~data:ignore
     |> ignore
   with Xmlm.Error ((l, c), e) as exn ->
     Printf.printf "%i %i %s\n" l c (error_message e);
     raise exn
 
 let () =
-  measure 100 "xmlm" xml_spec "xml" (fun () ->
-    parse xml_spec)
+  measure 100 "xmlm" "xml_spec" "xml" (fun () -> parse xml_spec)