perf: improve link label parsing performance

Author: lepture

src/mistune/core.py

@@ -133,6 +133,7 @@ def __init__(self, env: MutableMapping[str, Any]):
         self.no_close_bracket_before: int = 0  # high-water mark for DoS mitigation
         self.no_link_before: int = 0  # high-water mark for failed balanced link candidates
         self.link_brackets: Dict[int, Tuple[str, Dict[int, int]]] = {}
+        self.link_ranges: Dict[int, Tuple[str, List[int], List[int]]] = {}
 
     def prepend_token(self, token: Dict[str, Any]) -> None:
         """Insert token before the last token."""
@@ -149,6 +150,7 @@ def copy(self) -> "InlineState":
         state.image_depth = self.image_depth
         state.in_link = self.in_link
         state.link_brackets = self.link_brackets
+        state.link_ranges = self.link_ranges
         return state
 
 

src/mistune/inline_parser.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+from bisect import bisect_left
 from dataclasses import dataclass
 from typing import (
     Any,
@@ -142,26 +143,32 @@ def parse_link(self, m: Match[str], state: InlineState) -> Optional[int]:
             return pos
 
         text = None
+        text_start = pos
+        text_end = pos
         label, end_pos = parse_link_label(state.src, pos)
         if label is None:
             if pos <= state.no_close_bracket_before:
                 state.append_token({"type": "text", "raw": marker})
                 return pos
-            text, end_pos = _parse_link_text(state, pos)
-            if text is None:
-                if end_pos > state.no_close_bracket_before:
-                    state.no_close_bracket_before = end_pos
+            close_pos = _find_closing_bracket(state, pos)
+            if close_pos is None:
+                if len(state.src) > state.no_close_bracket_before:
+                    state.no_close_bracket_before = len(state.src)
                 return None
+            text_start = pos
+            text_end = close_pos
+            end_pos = close_pos + 1
 
         assert end_pos is not None
 
-        if text is None:
+        if label is not None:
             text = label
+            text_start = pos
+            text_end = end_pos - 1
 
-        assert text is not None
         body_end_pos = end_pos
 
-        has_nested_link = not is_image and self._contains_nested_link(text, state)
+        has_nested_link = not is_image and _label_contains_link(state, text_start, text_end)
         if has_nested_link:
             return None
         if end_pos >= len(state.src) and label is None:
@@ -179,6 +186,8 @@ def parse_link(self, m: Match[str], state: InlineState) -> Optional[int]:
                 # standard link [text](<url> "title")
                 attrs, pos2 = parse_link(state.src, end_pos + 1)
                 if pos2:
+                    if text is None:
+                        text = state.src[text_start:text_end]
                     token = self.__parse_link_token(is_image, text, attrs, state)
                     state.append_token(token)
                     return pos2
@@ -192,9 +201,13 @@ def parse_link(self, m: Match[str], state: InlineState) -> Optional[int]:
                         label = label2
 
         if label is None:
-            _mark_no_link_before(state, body_end_pos)
-            return None
-
+            ref_links = state.env.get("ref_links")
+            if not ref_links:
+                _mark_no_link_before(state, body_end_pos)
+                return None
+            if text is None:
+                text = state.src[text_start:text_end]
+            label = text
         ref_links = state.env.get("ref_links")
         if not ref_links:
             _mark_no_link_before(state, body_end_pos)
@@ -203,6 +216,8 @@ def parse_link(self, m: Match[str], state: InlineState) -> Optional[int]:
         key = unikey(label)
         env = ref_links.get(key)
         if env:
+            if text is None:
+                text = state.src[text_start:text_end]
             attrs = {"url": env["url"], "title": env.get("title")}
             token = self.__parse_link_token(is_image, text, attrs, state)
             token["ref"] = key
@@ -212,29 +227,6 @@ def parse_link(self, m: Match[str], state: InlineState) -> Optional[int]:
         _mark_no_link_before(state, body_end_pos)
         return None
 
-    def _contains_nested_link(self, text: str, state: InlineState) -> bool:
-        if "[" not in text:
-            return False
-        if "](" not in text and "][" not in text and not state.env.get("ref_links"):
-            return False
-
-        sc = self.compile_sc(["link"])
-        nested_state = state.copy()
-        nested_state.src = text
-        pos = 0
-        while pos < len(text):
-            m = sc.search(text, pos)
-            if not m:
-                return False
-
-            marker = m.group(0)
-            if marker == "[" and (m.start() == 0 or text[m.start() - 1] != "!"):
-                if _is_link_like(text, m.end(), nested_state):
-                    return True
-            pos = m.start() + 1
-
-        return False
-
     def __parse_link_token(
         self,
         is_image: bool,
@@ -849,56 +841,89 @@ def _merge_text_tokens(tokens: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     return result
 
 
-def _parse_link_text(state: InlineState, pos: int) -> Tuple[Optional[str], int]:
-    close_pos = _find_closing_bracket(state, pos)
-    if close_pos is None:
-        return None, len(state.src)
-    return state.src[pos:close_pos], close_pos + 1
-
-
 def _mark_no_link_before(state: InlineState, end_pos: int) -> None:
     if end_pos > state.no_link_before:
         state.no_link_before = end_pos
 
 
 def _find_closing_bracket(state: InlineState, pos: int) -> Optional[int]:
+    return _get_closing_bracket_map(state).get(pos)
+
+
+def _label_contains_link(state: InlineState, start: int, end: int) -> bool:
+    if start >= end:
+        return False
+
+    starts, suffix_min_ends = _get_link_range_index(state)
+    index = bisect_left(starts, start)
+    return index < len(starts) and starts[index] < end and suffix_min_ends[index] <= end
+
+
+def _get_link_range_index(state: InlineState) -> Tuple[List[int], List[int]]:
+    cache = state.link_ranges.get(id(state.src))
+    if cache is not None and cache[0] is state.src:
+        return cache[1], cache[2]
+
+    pairs = _get_closing_bracket_map(state)
+    ranges: List[Tuple[int, int]] = []
+    for label_start, close_pos in pairs.items():
+        opener = label_start - 1
+        if opener > 0 and state.src[opener - 1] == "!":
+            continue
+        link_end = _find_link_range_end(state.src, label_start, close_pos, state)
+        if link_end is not None:
+            ranges.append((opener, link_end))
+
+    ranges.sort()
+    starts = [start for start, _end in ranges]
+    suffix_min_ends = [0] * len(ranges)
+    min_end = len(state.src) + 1
+    for index in range(len(ranges) - 1, -1, -1):
+        end = ranges[index][1]
+        if end < min_end:
+            min_end = end
+        suffix_min_ends[index] = min_end
+
+    state.link_ranges[id(state.src)] = (state.src, starts, suffix_min_ends)
+    return starts, suffix_min_ends
+
+
+def _get_closing_bracket_map(state: InlineState) -> Dict[int, int]:
     cache = state.link_brackets.get(id(state.src))
     if cache is not None and cache[0] is state.src:
-        return cache[1].get(pos)
+        return cache[1]
 
     pairs = _build_closing_bracket_map(state.src)
     state.link_brackets[id(state.src)] = (state.src, pairs)
-    return pairs.get(pos)
-
+    return pairs
 
-def _is_link_like(src: str, pos: int, state: InlineState) -> bool:
-    label, end_pos = parse_link_label(src, pos)
-    if label is None:
-        label, end_pos = _parse_link_text(state, pos)
-        if label is None:
-            return False
 
-    assert label is not None
-    assert end_pos is not None
+def _find_link_range_end(src: str, label_start: int, close_pos: int, state: InlineState) -> Optional[int]:
+    end_pos = close_pos + 1
+    if end_pos < len(src):
+        marker = src[end_pos]
+        if marker == "(":
+            _attrs, new_pos = parse_link(src, end_pos + 1)
+            return new_pos
 
-    if end_pos >= len(src):
-        ref_links = state.env.get("ref_links")
-        return bool(ref_links and unikey(label) in ref_links)
-
-    marker = src[end_pos]
-    if marker == "(":
-        _attrs, new_pos = parse_link(src, end_pos + 1)
-        return bool(new_pos)
+        if marker == "[":
+            label, new_pos = parse_link_label(src, end_pos + 1)
+            if not new_pos:
+                return None
+            if label:
+                ref_label = label
+            else:
+                ref_label = src[label_start:close_pos]
 
-    if marker == "[":
-        label2, new_pos = parse_link_label(src, end_pos + 1)
-        if not new_pos:
-            return False
-        if label2:
-            label = label2
+            ref_links = state.env.get("ref_links")
+            if ref_links and unikey(ref_label) in ref_links:
+                return new_pos
+            return None
 
     ref_links = state.env.get("ref_links")
-    return bool(ref_links and unikey(label) in ref_links)
+    if ref_links and unikey(src[label_start:close_pos]) in ref_links:
+        return end_pos
+    return None
 
 
 def _build_closing_bracket_map(src: str) -> Dict[int, int]:

tests/test_security_inline.py

@@ -113,3 +113,19 @@ def test_deep_link_input_does_not_recurse_forever(self):
         elapsed = time.monotonic() - start
 
         self.assertLess(elapsed, DEADLINE)
+
+    def test_repeated_link_suffixes_are_near_linear(self):
+        md = create_markdown()
+
+        def render(size):
+            text = "[" * size + "a" + "](/u)" * size + "\n"
+            start = time.process_time()
+            md(text)
+            return time.process_time() - start
+
+        render(200)
+        small = render(1000)
+        large = render(2000)
+
+        self.assertLess(large, DEADLINE)
+        self.assertLess(large, small * 3)