Make memo tag determination based on memoized component (#6453)

* Make memo tag determination based on memoized component

Calculate the memo component's hash, not the hash of the pre-memoized component.

This allows PASSTHROUGH children to memoize to the same underlying function and avoid a recursive .render() call.

Updating the `_get_component_hash` function to hash the imports, hooks, custom
code, and app wraps ensures that structurally similar component do not collapse
to the same memo component if critical aspects differ. Previously the hash was
only based on the rendered component itself, so hooks, like on_mount would not
be factored in.

* avoid recursive calls for PASSTHROUGH memo components

* fix _get_all_refs for form on_submit

Author: masenf

packages/reflex-base/src/reflex_base/components/component.py

@@ -1381,6 +1381,71 @@ def render(self) -> dict:
         self._cached_render_result = rendered_dict
         return rendered_dict
 
+    def _get_component_hash(self, shallow: bool = False) -> str:
+        """Get a stable content hash for this component.
+
+        The hash incorporates the rendered JSX dict plus the component's
+        recursive imports, hooks (including internal lifecycle hooks),
+        custom code, and app-wrap components, so two components that
+        compile to semantically distinct JS modules hash differently
+        even when their ``render()`` output happens to match (e.g. two
+        components differing only in ``on_mount``, which is excluded
+        from ``_render`` props but lives in the lifecycle hook).
+
+        Args:
+            shallow: If True, only hash the component's own render output and
+                directly defined hooks, imports, custom code, and app-wrap
+                components, excluding any of those from child components.
+
+        Returns:
+            The hex digest content hash.
+        """
+        hasher = md5(usedforsecurity=False)
+        _update_deterministic_hash(hasher, self.render())
+        if shallow:
+            # For non-snapshot strategies, we only hash the component's own hooks, imports, custom code, and app-wrap components
+            _update_deterministic_hash(hasher, dict(self._get_imports()))
+            _update_deterministic_hash(hasher, dict(self._get_hooks_internal()))
+            _update_deterministic_hash(hasher, dict(self._get_added_hooks()))
+            _update_deterministic_hash(hasher, self._get_hooks())
+            _update_deterministic_hash(hasher, self._get_custom_code())
+            _update_deterministic_hash(hasher, dict(self._get_app_wrap_components()))
+        else:
+            _update_deterministic_hash(hasher, dict(self._get_all_imports()))
+            _update_deterministic_hash(hasher, dict(self._get_all_hooks_internal()))
+            _update_deterministic_hash(hasher, dict(self._get_all_hooks()))
+            _update_deterministic_hash(hasher, dict(self._get_all_custom_code()))
+            _update_deterministic_hash(
+                hasher, dict(self._get_all_app_wrap_components())
+            )
+        return hasher.hexdigest()
+
+    def _compute_memo_tag(self) -> str:
+        """Compute a stable tag name for memoizing this component.
+
+        The class qualname is encoded directly in the tag prefix so that
+        distinct classes which happen to render identically never collide
+        on a tag. Tag collision would silently share a single cached memo
+        wrapper across classes and drop the later class's class-level
+        metadata (e.g. ``_get_app_wrap_components``, which carries
+        providers like ``UploadFilesProvider`` that must reach the app
+        root).
+
+        Returns:
+            The stable tag name.
+        """
+        from reflex_base.components.memoize_helpers import (
+            MemoizationStrategy,
+            get_memoization_strategy,
+        )
+
+        comp_hash = self._get_component_hash(
+            shallow=get_memoization_strategy(self) == MemoizationStrategy.PASSTHROUGH
+        )
+        return format.format_state_name(
+            f"{type(self).__qualname__}_{self.tag or 'Comp'}_{comp_hash}"
+        ).capitalize()
+
     def _replace_prop_names(self, rendered_dict: dict) -> None:
         """Replace the prop names in the render dictionary.
 

pyi_hashes.json

@@ -120,5 +120,5 @@
   "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "2c5fadcc014056f041cd4d916137d9e7",
   "reflex/__init__.pyi": "3a9bb8544cbc338ffaf0a5927d9156df",
   "reflex/components/__init__.pyi": "f39a2af77f438fa243c58c965f19d42e",
-  "reflex/experimental/memo.pyi": "9946d9b757f7cef5f53d599194d6e50e"
+  "reflex/experimental/memo.pyi": "82d8699470071df80886a4a6ba8dccfe"
 }

reflex/compiler/plugins/memoize.py

@@ -22,12 +22,7 @@
 import dataclasses
 from typing import Any
 
-from reflex_base.components.component import (
-    BaseComponent,
-    Component,
-    _deterministic_hash,
-    _hash_str,
-)
+from reflex_base.components.component import BaseComponent, Component
 from reflex_base.components.memoize_helpers import (
     MemoizationStrategy,
     _is_structural_memoization_child,
@@ -38,40 +33,10 @@
 from reflex_base.constants.compiler import MemoizationDisposition
 from reflex_base.plugins import ComponentAndChildren, PageContext
 from reflex_base.plugins.base import Plugin
-from reflex_base.utils import format
 
 from reflex.experimental.memo import create_passthrough_component_memo
 
 
-def _compute_memo_tag(component: Component) -> str | None:
-    """Compute a stable tag name for a memoizable component.
-
-    Returns ``None`` for components that render empty (non-visual components
-    are never memoized).
-
-    The class qualname is encoded directly in the tag prefix so that distinct
-    classes which render identically never collide on a tag. Tag collision
-    would silently share a single cached memo wrapper across classes and drop
-    the later class's class-level metadata (e.g. ``_get_app_wrap_components``,
-    which carries providers like ``UploadFilesProvider`` that must reach the
-    app root). Baking the qualname into the prefix avoids re-concatenating
-    the rendered JSX into the hash input on every call.
-
-    Args:
-        component: The component to name.
-
-    Returns:
-        The stable tag name, or ``None`` if the component renders empty.
-    """
-    rendered_code = component.render()
-    if not rendered_code:
-        return None
-    code_hash = _hash_str(_deterministic_hash(rendered_code))
-    return format.format_state_name(
-        f"{type(component).__qualname__}_{component.tag or 'Comp'}_{code_hash}"
-    ).capitalize()
-
-
 def _subtree_has_reactive_data(
     component: Component, _cache: dict[int, bool] | None = None
 ) -> bool:
@@ -383,16 +348,17 @@ def _build_wrapper(
             The wrapper instance, or ``None`` if the component's render is
             empty and has no meaningful tag.
         """
-        tag = _compute_memo_tag(comp)
-        if tag is None:
-            return None
-
         comp = fix_event_triggers_for_memo(comp, page_context)
 
-        compile_context.memoize_wrappers[tag] = None
         # Passthrough memo definitions capture app-specific event/state vars, so
-        # they must be rebuilt for each compile instead of shared globally.
-        wrapper_factory, definition = create_passthrough_component_memo(tag, comp)
+        # they must be rebuilt for each compile instead of shared globally. The
+        # tag is derived from the rendered memo body inside
+        # ``create_passthrough_component_memo`` after the ``{children}`` hole
+        # is substituted, so passthrough wrappers that differ only in their
+        # children collapse to a single definition.
+        wrapper_factory, definition = create_passthrough_component_memo(comp)
+        tag = definition.export_name
+        compile_context.memoize_wrappers[tag] = None
         compile_context.auto_memo_components[tag] = definition
 
         wrapper = wrapper_factory()

reflex/experimental/memo.py

@@ -1008,7 +1008,6 @@ def _create_component_wrapper(
 
 
 def create_passthrough_component_memo(
-    export_name: str,
     component: Component,
 ) -> tuple[
     Callable[..., ExperimentalMemoComponent],
@@ -1020,8 +1019,13 @@ def create_passthrough_component_memo(
     through the experimental memo pipeline instead of emitting ad-hoc page-local
     ``React.memo`` declarations.
 
+    The exported memo name is derived from ``component._compute_memo_tag()``
+    after the ``{children}`` hole has been substituted into the wrapped
+    component's children (passthrough mode), so two call-sites differing only
+    in their children — whose generated memo bodies are identical — collapse
+    to one wrapper.
+
     Args:
-        export_name: The exported memo component name.
         component: The component to wrap.
 
     Returns:
@@ -1044,22 +1048,46 @@ def passthrough(children: Var[Component]) -> Component:
         new_component = copy(component)
         if render_snapshot:
             return new_component
-        # Keep ``new_component.children`` as the ORIGINAL children so
-        # compile-time walkers that introspect the subtree (e.g. Form's
-        # ``_get_form_refs``) see the real descendants. The ``{children}``
-        # hole lives on the definition and the compiler swaps it in only for
-        # JSX render / imports collection.
-        captured_hole_child.append(Bare.create(children))
+        hole_bare = Bare.create(children)
+        captured_hole_child.append(hole_bare)
+        # Substitute the ``{children}`` hole for the original descendants so
+        # the memo body's hash and JSX both reflect the placeholder, not the
+        # specific children at any given call site. Original descendants stay
+        # reachable on the page-level wrapper via the plugin's
+        # ``_get_all_refs`` delegation back to the source component.
+        new_component.children = [hole_bare]
+        # Compile-time walkers that need the real subtree (notably
+        # ``Form._get_form_refs`` collecting id-based input refs into the
+        # generated ``handleSubmit`` JS) call ``self._get_all_refs()`` while
+        # the memo body's hooks are computed. With the hole substituted in,
+        # that walk would return nothing and the form handler would emit an
+        # empty ``field_ref_mapping``. Delegate ref collection back to the
+        # source component so descendants behind the hole remain visible.
+        object.__setattr__(new_component, "_get_all_refs", component._get_all_refs)
         return new_component
 
-    passthrough.__name__ = format.to_snake_case(export_name)
+    # Evaluate once to compute the tag from the rendered memo body shape.
+    # ``_create_component_definition`` will evaluate again internally; the
+    # second pass overwrites ``captured_hole_child`` but the captured value
+    # is identical.
+    params = _analyze_params(passthrough, for_component=True)
+    preview = _normalize_component_return(_evaluate_memo_function(passthrough, params))
+    if preview is None:
+        msg = (
+            "`create_passthrough_component_memo` requires a component that "
+            "normalizes to `rx.Component`."
+        )
+        raise TypeError(msg)
+    tag = preview._compute_memo_tag()
+
+    passthrough.__name__ = format.to_snake_case(tag)
     passthrough.__qualname__ = passthrough.__name__
     passthrough.__module__ = __name__
 
     definition = _create_component_definition(passthrough, Component)
     replacements: dict[str, Any] = {}
-    if definition.export_name != export_name:
-        replacements["export_name"] = export_name
+    if definition.export_name != tag:
+        replacements["export_name"] = tag
     if captured_hole_child:
         replacements["passthrough_hole_child"] = captured_hole_child[0]
     if replacements:

tests/units/compiler/test_memoize_plugin.py

@@ -196,6 +196,50 @@ def test_memoize_wrapper_deduped_across_repeated_subtrees() -> None:
     ) == 1
 
 
+def test_memoize_wrappers_distinct_for_different_on_mount() -> None:
+    """Two components differing only in ``on_mount`` must NOT dedupe.
+
+    ``on_mount`` is excluded from ``_render``'s props, so its handler does not
+    appear in ``component.render()``. The memo wrapper body, however, includes
+    a ``useEffect`` lifecycle hook whose body invokes the ``on_mount`` handler
+    — two distinct handlers produce two distinct memo bodies and must compile
+    to two distinct memo wrappers. Hashing only ``render()`` lets them collide
+    on a single tag, silently dropping one handler's logic.
+    """
+    ctx, _page_ctx = _compile_single_page(
+        lambda: Fragment.create(
+            Plain.create(on_mount=rx.console_log("a")),
+            Plain.create(on_mount=rx.console_log("b")),
+        )
+    )
+    assert len(ctx.memoize_wrappers) == 2, (
+        "Components with different on_mount handlers must produce distinct "
+        f"memo wrappers, got: {list(ctx.memoize_wrappers)}"
+    )
+
+
+def test_memoize_wrappers_dedupe_passthrough_with_different_children() -> None:
+    """Passthrough memos with identical props but different children must dedupe.
+
+    Passthrough memo bodies render a ``{children}`` placeholder rather than
+    the captured child JSX — the actual children flow through at the page-side
+    call site. Two components with identical props but different children
+    therefore generate identical memo body code and must share one wrapper
+    tag. Hashing the rendered children into the tag splits them needlessly,
+    bloating the generated component module with duplicate definitions.
+    """
+    ctx, _page_ctx = _compile_single_page(
+        lambda: Fragment.create(
+            WithProp.create("apple", label=STATE_VAR),
+            WithProp.create("banana", label=STATE_VAR),
+        )
+    )
+    assert len(ctx.memoize_wrappers) == 1, (
+        "Passthrough memos differing only in children must collapse to one "
+        f"wrapper, got: {list(ctx.memoize_wrappers)}"
+    )
+
+
 @pytest.mark.parametrize(
     ("special_form", "body_marker"),
     [
@@ -367,9 +411,7 @@ def test_common_memoization_snapshot_helper_classifies_snapshot_cases() -> None:
 
 def test_generated_memo_component_is_not_itself_memoized() -> None:
     """The generated memo component instance itself is skipped by the heuristic."""
-    wrapper_factory, _definition = create_passthrough_component_memo(
-        "MyTag", Fragment.create()
-    )
+    wrapper_factory, _definition = create_passthrough_component_memo(Fragment.create())
     wrapper = wrapper_factory(Plain.create())
     assert isinstance(wrapper, ExperimentalMemoComponent)
     assert not _should_memoize(wrapper)
@@ -400,14 +442,16 @@ def test_event_trigger_memoization_not_emit_usecallback_in_page_hooks() -> None:
 
 def test_generated_memo_component_renders_as_its_exported_tag() -> None:
     """The generated experimental memo component renders as its exported tag."""
-    wrapper_factory, definition = create_passthrough_component_memo(
-        "MyWrapper_abc", Fragment.create()
-    )
+    wrapper_factory, definition = create_passthrough_component_memo(Fragment.create())
     wrapper = wrapper_factory(Plain.create())
     assert isinstance(wrapper, ExperimentalMemoComponent)
-    assert wrapper.tag == "MyWrapper_abc"
-    assert definition.export_name == "MyWrapper_abc"
-    assert wrapper.render()["name"] == "MyWrapper_abc"
+    tag = definition.export_name
+    assert tag.startswith("Fragment_"), (
+        f"Expected the wrapped class qualname to be encoded in the tag prefix; "
+        f"got {tag!r}"
+    )
+    assert wrapper.tag == tag
+    assert wrapper.render()["name"] == tag
 
 
 def test_passthrough_memo_definitions_are_not_shared_globally(monkeypatch) -> None:
@@ -421,18 +465,14 @@ def test_passthrough_memo_definitions_are_not_shared_globally(monkeypatch) -> No
     first_component = Plain.create(STATE_VAR)
     second_component = Plain.create(STATE_VAR)
 
-    monkeypatch.setattr(memoize_plugin, "_compute_memo_tag", lambda comp: tag)
     monkeypatch.setattr(
         memoize_plugin,
         "fix_event_triggers_for_memo",
         lambda comp, page_context: comp,
     )
 
-    def fake_create_passthrough_component_memo(
-        export_name: str,
-        component: Component,
-    ):
-        definition = SimpleNamespace(export_name=export_name, component=component)
+    def fake_create_passthrough_component_memo(component: Component):
+        definition = SimpleNamespace(export_name=tag, component=component)
         return (lambda definition=definition: definition), definition
 
     monkeypatch.setattr(