Auto-memoize user @rx.memo components with stateful props

docs/library/other/memo.md

@@ -68,6 +68,18 @@ def index():
     )
 ```
 
+Binding state to a prop at the call site does not pull the state into the page.
+The compiler moves that call into a generated wrapper component that holds the
+state hooks the prop needs, so the page itself keeps no dependency on the state.
+When the state changes, the wrapper re-renders and React's `memo` stops there
+unless the prop's value actually changed. The page function itself never re-runs,
+so nothing in it re-renders except the components that read the changed state
+themselves — each inside its own wrapper, the `rx.input` above included.
+
+That makes the call site the place to punch a single dependency through to an
+expensive component: pass exactly the Vars it needs, and it re-renders for those
+and nothing else, however much the rest of the state churns.
+
 ## Using with `rx.foreach`
 
 To render a memoized component for each item of a list Var, wrap the call in a

news/6949.performance.md

@@ -0,0 +1 @@
+`@rx.memo` components with props bound to state are now auto-memoized at the call site: the state hooks (and event-handler callbacks) those props need compile into a generated wrapper component instead of the page module. A state change re-renders that wrapper rather than the whole page, and React's `memo` stops there unless one of the prop values actually changed — so binding a Var at the call site scopes an expensive component to exactly the state it reads, instead of coupling it to the page.

packages/reflex-base/news/6949.performance.md

@@ -0,0 +1 @@
+`MemoComponent` instances no longer opt out of compiler auto-memoization wholesale. Only the passthrough wrappers the auto-memoize pass generates do, tracked by the new `auto_memo_wrapper` flag on `MemoComponentDefinition`, so state-bound props and event handlers on a `@rx.memo` call site compile their hooks into a generated wrapper instead of the enclosing page.

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

@@ -324,6 +324,10 @@ class MemoComponentDefinition(MemoDefinition):
     # wrapper's ``VarData`` supplies its imports, so a custom wrapper brings
     # its own and ``None`` pulls in nothing.
     wrapper: Var | None = DEFAULT_MEMO_WRAPPER
+    # Set for definitions the compiler's auto-memoize pass creates (see
+    # ``create_passthrough_component_memo``). Instances of such a definition
+    # are the auto-memo boundary itself, so the pass must not wrap them again.
+    auto_memo_wrapper: bool = False
     # The name React DevTools shows for this memo. ``export_name`` (derived
     # from the decorated function) is already readable for ``@rx.memo``, but
     # auto-memoized wrappers carry a hash-suffixed tag, so the plugin sets this
@@ -341,10 +345,23 @@ def component(self) -> Component:
 
 
 class MemoComponent(Component):
-    """A rendered instance of a memo component."""
+    """A rendered instance of a memo component.
+
+    Instances take part in compiler auto-memoization like any other component.
+    A call site binding state Vars (or event handlers) to props *must* be
+    wrapped, so those hooks compile into the generated wrapper instead of the
+    page module: otherwise every state change re-renders the whole page, and
+    React's ``memo`` on this component only spares its own subtree. With the
+    wrapper in place, the page holds no state hook, the wrapper absorbs the
+    re-render, and this component re-renders only when a bound prop value
+    actually changes.
+
+    Wrappers the auto-memoize pass generates are themselves ``MemoComponent``
+    instances; they opt out via ``MemoizationDisposition.NEVER`` (see
+    :func:`_get_memo_component_class`) since they already are the boundary.
+    """
 
     library = f"$/{constants.Dirs.COMPONENTS_PATH}"
-    _memoization_mode = MemoizationMode(disposition=MemoizationDisposition.NEVER)
 
     # The user-authored component class this wrapper stands in for. Populated
     # on the dynamic subclass by ``_get_memo_component_class`` so
@@ -390,6 +407,7 @@ def _get_memo_component_class(
     export_name: str,
     wrapped_component_type: type[Component] = Component,
     source_module: str | None = None,
+    auto_memo_wrapper: bool = False,
 ) -> type[MemoComponent]:
     """Get the component subclass for a memo export.
 
@@ -407,6 +425,11 @@ def _get_memo_component_class(
         source_module: The user-app Python module that defined this memo. When
             set, the wrapper imports from a path mirroring that module instead
             of the per-name ``utils/components/<name>`` path.
+        auto_memo_wrapper: Whether the export is a wrapper generated by the
+            compiler's auto-memoize pass. Such wrappers already are the memo
+            boundary, so they opt out of being auto-memoized themselves;
+            user-authored ``@rx.memo`` components do not, so their stateful
+            props land in a generated wrapper instead of the page module.
 
     Returns:
         A cached component subclass with the tag set at class definition time.
@@ -421,6 +444,10 @@ def _get_memo_component_class(
         "library": library,
         "_wrapped_component_type": wrapped_component_type,
     }
+    if auto_memo_wrapper:
+        attrs["_memoization_mode"] = MemoizationMode(
+            disposition=MemoizationDisposition.NEVER
+        )
     if (
         wrapped_component_type._get_app_wrap_components
         is not Component._get_app_wrap_components
@@ -1717,6 +1744,7 @@ def __call__(self, *children: Any, **props: Any) -> MemoComponent:
             definition.export_name,
             type(component),
             definition.source_module,
+            definition.auto_memo_wrapper,
         )._create(
             children=list(children),
             memo_definition=definition,
@@ -1856,7 +1884,10 @@ def passthrough(children: Var[Component]) -> Component:
     definition = _create_component_definition(passthrough, Component, source_module)
     # ``export_name`` is the content-hashed tag, which reads as noise in the
     # React DevTools tree. Name the memo after the Python class it wraps.
-    replacements: dict[str, Any] = {"display_name": type(component).__qualname__}
+    replacements: dict[str, Any] = {
+        "auto_memo_wrapper": True,
+        "display_name": type(component).__qualname__,
+    }
     if definition.export_name != tag:
         replacements["export_name"] = tag
     if captured_hole_child:

pyi_hashes.json

@@ -120,5 +120,5 @@
   "packages/reflex-components-sonner/src/reflex_components_sonner/toast.pyi": "f170ac685b6ba5892370166c80684db3",
   "reflex/__init__.pyi": "a3e1782fab4a9aed55f66cc98af8c217",
   "reflex/components/__init__.pyi": "9facd05a776d0641432696bbf8e34388",
-  "reflex/experimental/memo.pyi": "35583b85befadf5cb125b14f7cd459cb"
+  "reflex/experimental/memo.pyi": "7bfdf4841052a4b8d0df753d4184ef4b"
 }

reflex/compiler/plugins/memoize.py

@@ -134,6 +134,12 @@ def _should_memoize(component: Component) -> bool:
     are evaluated from their own props/triggers; descendants are visited
     independently by the walker.
 
+    Explicitly memoized (``@rx.memo``) components are no exception: React's
+    ``memo`` only spares their own subtree, so state bound at the call site
+    still needs a wrapper to keep the hooks out of the page module. The
+    wrappers this pass generates are themselves memo components and opt out
+    via ``MemoizationDisposition.NEVER``.
+
     Args:
         component: The candidate component.
 

tests/integration/tests_playwright/test_memo.py

@@ -81,6 +81,16 @@ def keyed_row(label: rx.Var[str]) -> rx.Component:
         # element id so each row is locatable after reordering.
         return rx.input(id=label)
 
+    @rx.memo
+    def framed(title: rx.Var[str], children: rx.Var[rx.Component]) -> rx.Component:
+        # Stateful prop *and* a children slot: the auto-memoize pass wraps the
+        # call site so the state hooks live in the generated wrapper, which
+        # passes the page-rendered children straight through.
+        return rx.vstack(
+            rx.text(title, id="framed-title"),
+            rx.box(children, id="framed-slot"),
+        )
+
     @rx.memo(wrapper=None)
     def unwrapped_label(value: rx.Var[str]) -> rx.Component:
         # Compiled without the React ``memo`` wrapper: a bare function
@@ -112,6 +122,10 @@ def index() -> rx.Component:
                 id="keyed-rows",
             ),
             unwrapped_label(value=MemoState.last_value),
+            framed(
+                rx.text(MemoState.last_value, id="framed-child"),
+                title=MemoState.last_value,
+            ),
         )
 
     app = rx.App()
@@ -246,6 +260,32 @@ def test_memo_key_preserves_identity_across_reorder(
         expect(page.locator(f"#{row_id}")).to_have_value(row_id.upper())
 
 
+def test_memo_stateful_prop_and_children_update(
+    memo_app: AppHarness, page: Page
+) -> None:
+    """A memo bound to state renders its children and follows state changes.
+
+    The call site binds a state Var to a prop and passes children positionally,
+    so the auto-memoize pass hoists the state hooks into a generated wrapper
+    that feeds both the prop and the page-rendered children.
+
+    Args:
+        memo_app: Running app harness.
+        page: Playwright page.
+    """
+    _load_page(page, memo_app)
+
+    expect(page.locator("#framed-title")).to_have_text("")
+    expect(page.locator("#framed-child")).to_have_text("")
+
+    page.locator("#memo-input").fill("framed_update")
+
+    expect(page.locator("#framed-title")).to_have_text("framed_update")
+    expect(page.locator("#framed-slot").locator("#framed-child")).to_have_text(
+        "framed_update"
+    )
+
+
 def test_memo_wrapper_none_renders_and_updates(
     memo_app: AppHarness, page: Page
 ) -> None: