perf: improve performance for math and formatting plugins

Author: lepture

src/mistune/core.py

@@ -134,6 +134,7 @@ def __init__(self, env: MutableMapping[str, Any]):
         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]]] = {}
+        self.formatting_no_end: Dict[Tuple[int, str], Tuple[str, int]] = {}
 
     def prepend_token(self, token: Dict[str, Any]) -> None:
         """Insert token before the last token."""
@@ -151,6 +152,7 @@ def copy(self) -> "InlineState":
         state.in_link = self.in_link
         state.link_brackets = self.link_brackets
         state.link_ranges = self.link_ranges
+        state.formatting_no_end = self.formatting_no_end
         return state
 
 

src/mistune/plugins/formatting.py

@@ -61,8 +61,14 @@ def _parse_to_end(
     marker: str,
 ) -> Optional[int]:
     pos = m.end()
+    cache_key = (id(state.src), marker)
+    cache = state.formatting_no_end.get(cache_key)
+    if cache is not None and cache[0] is state.src and pos <= cache[1]:
+        return None
+
     end_pos = _find_end_marker(state.src, pos, marker)
     if end_pos is None:
+        state.formatting_no_end[cache_key] = (state.src, len(state.src))
         return None
     text = state.src[pos : end_pos - 2]
     new_state = state.copy()

tests/test_security_formatting.py

@@ -34,3 +34,25 @@ def test_repeated_formatting_pairs_return_quickly(self):
                 elapsed = time.monotonic() - start
 
                 self.assertLess(elapsed, deadline)
+
+    def test_repeated_unclosed_formatting_markers_are_near_linear(self):
+        cases = [
+            ("strikethrough", "~~a "),
+            ("mark", "==a "),
+            ("insert", "^^a "),
+        ]
+        for plugin, unit in cases:
+            with self.subTest(plugin=plugin):
+                md = create_markdown(plugins=[plugin])
+
+                def render(size):
+                    start = time.process_time()
+                    md(unit * size + "\n")
+                    return time.process_time() - start
+
+                render(200)
+                small = render(1000)
+                large = render(2000)
+
+                self.assertLess(large, 1.0 if IS_PYPY else 0.5)
+                self.assertLess(large, small * 3)

tests/test_security_math.py

@@ -1,7 +1,11 @@
+import platform
+import time
 from unittest import TestCase
 
 from mistune import create_markdown
 
+IS_PYPY = platform.python_implementation() == "PyPy"
+
 
 class TestMathSecurity(TestCase):
     def test_math_plugin_escapes_math_content(self):
@@ -14,3 +18,18 @@ def test_math_plugin_escapes_math_content(self):
         self.assertIn("<script>alert(1)</script>", block_html)
         self.assertNotIn("<script>", inline_html)
         self.assertNotIn("<script>", block_html)
+
+    def test_repeated_inline_math_openers_are_near_linear(self):
+        md = create_markdown(plugins=["math"])
+
+        def render(size):
+            start = time.process_time()
+            md("$a " * size + "\n")
+            return time.process_time() - start
+
+        render(200)
+        small = render(1000)
+        large = render(2000)
+
+        self.assertLess(large, 1.0 if IS_PYPY else 0.5)
+        self.assertLess(large, small * 3)