From 0eb568be38997536d5868023a6bfa70d86069d7b Mon Sep 17 00:00:00 2001
From: Masen Furer <m_github@0x26.net>
Date: Thu, 30 Nov 2023 16:41:25 -0800
Subject: [PATCH] [REF-1352] Markdown component_map hash improvements

* Hash against the compiled form of each component, not the function that compiles it
* LRU cache the `get_base_component_map` function, since that always returns the
  same thing, no need to recreate the map for every instance of every markdown.
* Only hash the component map once at create() time
---
 reflex/components/typography/markdown.py  | 24 ++++++++++++++++++-----
 reflex/components/typography/markdown.pyi |  4 ++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/reflex/components/typography/markdown.py b/reflex/components/typography/markdown.py
index df6154d59f..fc7f84a212 100644
--- a/reflex/components/typography/markdown.py
+++ b/reflex/components/typography/markdown.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import textwrap
+from functools import lru_cache
 from hashlib import md5
 from typing import Any, Callable, Dict, Union
 
@@ -35,6 +36,7 @@
 
 
 # Component Mapping
+@lru_cache
 def get_base_component_map() -> dict[str, Callable]:
     """Get the base component map.
 
@@ -89,6 +91,9 @@ class Markdown(Component):
     # Custom styles for the markdown (deprecated in v0.2.9).
     custom_styles: Dict[str, Any] = {}
 
+    # The hash of the component map, generated at create() time.
+    component_map_hash: str = ""
+
     @classmethod
     def create(cls, *children, **props) -> Component:
         """Create a markdown component.
@@ -124,7 +129,12 @@ def create(cls, *children, **props) -> Component:
             src = textwrap.dedent(src)
 
         # Create the component.
-        return super().create(src, component_map=component_map, **props)
+        return super().create(
+            src,
+            component_map=component_map,
+            component_map_hash=cls._component_map_hash(component_map),
+            **props,
+        )
 
     def get_custom_components(
         self, seen: set[str] | None = None
@@ -264,11 +274,15 @@ def format_component_map(self) -> dict[str, str]:
 
         return components
 
-    def _component_map_hash(self) -> str:
-        return md5(str(self.component_map).encode()).hexdigest()
+    @staticmethod
+    def _component_map_hash(component_map) -> str:
+        inp = str(
+            {tag: component(_MOCK_ARG) for tag, component in component_map.items()}
+        ).encode()
+        return md5(inp).hexdigest()
 
     def _get_component_map_name(self) -> str:
-        return f"ComponentMap_{self._component_map_hash()}"
+        return f"ComponentMap_{self.component_map_hash}"
 
     def _get_custom_code(self) -> str | None:
         hooks = set()
@@ -292,7 +306,7 @@ def _render(self) -> Tag:
                 remark_plugins=_REMARK_PLUGINS,
                 rehype_plugins=_REHYPE_PLUGINS,
             )
-            .remove_props("componentMap")
+            .remove_props("componentMap", "componentMapHash")
         )
         tag.special_props.add(
             Var.create_safe(
diff --git a/reflex/components/typography/markdown.pyi b/reflex/components/typography/markdown.pyi
index 3a0fb43a88..995b96d749 100644
--- a/reflex/components/typography/markdown.pyi
+++ b/reflex/components/typography/markdown.pyi
@@ -8,6 +8,7 @@ from reflex.vars import Var, BaseVar, ComputedVar
 from reflex.event import EventChain, EventHandler, EventSpec
 from reflex.style import Style
 import textwrap
+from functools import lru_cache
 from hashlib import md5
 from typing import Any, Callable, Dict, Union
 from reflex.compiler import utils
@@ -32,6 +33,7 @@ _REHYPE_KATEX = Var.create_safe("rehypeKatex", _var_is_local=False)
 _REHYPE_RAW = Var.create_safe("rehypeRaw", _var_is_local=False)
 _REHYPE_PLUGINS = Var.create_safe([_REHYPE_KATEX, _REHYPE_RAW])
 
+@lru_cache
 def get_base_component_map() -> dict[str, Callable]: ...
 
 class Markdown(Component):
@@ -42,6 +44,7 @@ class Markdown(Component):
         *children,
         component_map: Optional[Dict[str, Any]] = None,
         custom_styles: Optional[Dict[str, Any]] = None,
+        component_map_hash: Optional[str] = None,
         style: Optional[Style] = None,
         key: Optional[Any] = None,
         id: Optional[Any] = None,
@@ -101,6 +104,7 @@ class Markdown(Component):
             *children: The children of the component.
             component_map: The component map from a tag to a lambda that creates a component.
             custom_styles: Custom styles for the markdown (deprecated in v0.2.9).
+            component_map_hash: The hash of the component map, generated at create() time.
             style: The style of the component.
             key: A unique key for the component.
             id: The id for the component.