diff --git a/docs/library/other/memo.md b/docs/library/other/memo.md index 25a64852571..9786d059f7d 100644 --- a/docs/library/other/memo.md +++ b/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 diff --git a/news/6949.performance.md b/news/6949.performance.md new file mode 100644 index 00000000000..4aadd08a79c --- /dev/null +++ b/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. diff --git a/packages/reflex-base/news/6949.performance.md b/packages/reflex-base/news/6949.performance.md new file mode 100644 index 00000000000..35130a73f45 --- /dev/null +++ b/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. diff --git a/packages/reflex-base/src/reflex_base/components/memo.py b/packages/reflex-base/src/reflex_base/components/memo.py index 71e874a1437..a00842540fd 100644 --- a/packages/reflex-base/src/reflex_base/components/memo.py +++ b/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/`` 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: diff --git a/pyi_hashes.json b/pyi_hashes.json index 172d76bf7d4..e089f6ed716 100644 --- a/pyi_hashes.json +++ b/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" } diff --git a/reflex/compiler/plugins/memoize.py b/reflex/compiler/plugins/memoize.py index 36aee008ae7..a50e3e30569 100644 --- a/reflex/compiler/plugins/memoize.py +++ b/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. diff --git a/tests/integration/tests_playwright/test_memo.py b/tests/integration/tests_playwright/test_memo.py index 9aadeda53bc..1c632a50652 100644 --- a/tests/integration/tests_playwright/test_memo.py +++ b/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: diff --git a/tests/units/compiler/test_memoize_plugin.py b/tests/units/compiler/test_memoize_plugin.py index 743536857a2..0098e49bd76 100644 --- a/tests/units/compiler/test_memoize_plugin.py +++ b/tests/units/compiler/test_memoize_plugin.py @@ -127,6 +127,12 @@ class SpecialFormMemoState(BaseState): value: Field[str] = field(default="a") +class MemoTriggerState(BaseState): + @rx.event + def ping(self): + """No-op handler for event-trigger memoization tests.""" + + @dataclasses.dataclass(slots=True) class FakePage: route: str @@ -547,6 +553,197 @@ def test_generated_memo_component_is_not_itself_memoized() -> None: assert not _should_memoize(wrapper) +def test_auto_memo_wrapper_opts_out_of_being_memoized() -> None: + """Generated wrappers carry ``NEVER`` so the pass can't wrap them again. + + The wrapper is itself a ``MemoComponent``; without the opt-out, a wrapper + built around a stateful component would look eligible to the heuristic and + the pass would wrap wrappers forever. + """ + from reflex_base.event import EventChain + + wrapper_factory, definition = create_passthrough_component_memo( + WithProp.create(label=STATE_VAR) + ) + assert definition.auto_memo_wrapper + wrapper = wrapper_factory() + assert isinstance(wrapper, MemoComponent) + assert wrapper._memoization_mode.disposition is MemoizationDisposition.NEVER + assert not _should_memoize(wrapper) + + # Even a signal the heuristic normally treats as eligible must not win. + wrapper.event_triggers["on_click"] = Var(_js_expr="test_event")._replace( + _var_type=EventChain, + merge_var_data=VarData(state="TestState"), + ) + assert not _should_memoize(wrapper) + + +def test_user_memo_with_stateful_prop_is_auto_memoized() -> None: + """An ``@rx.memo`` component bound to state gets its own memo wrapper. + + Regression: ``MemoComponent`` used to opt out of auto-memoization + wholesale, so binding a state Var at the call site left the state + ``useContext`` in the page module — every state change then re-rendered + the whole page, including static siblings. The hooks must live in a + generated wrapper instead, which re-renders on state change and lets + React's ``memo`` skip the wrapped component unless a prop value changed. + """ + from reflex.compiler.compiler import compile_memo_components + + @rx.memo + def stateful_card(label: rx.Var[str]) -> Component: + return WithProp.create(label=label) + + ctx, page_ctx = _compile_single_page( + lambda: Fragment.create( + Plain.create(LiteralVar.create("static sibling")), + stateful_card(label=SpecialFormMemoState.value), + ) + ) + + page_output = page_ctx.output_code + assert page_output is not None + assert "useContext(StateContexts" not in page_output + assert not any("useContext(StateContexts" in hook for hook in page_ctx.hooks) + + (definition,) = ctx.auto_memo_components.values() + assert isinstance(definition, MemoComponentDefinition) + wrapped = definition.component + assert isinstance(wrapped, MemoComponent) + assert wrapped.tag is not None + assert wrapped.tag.startswith("StatefulCard") + + # The page renders the wrapper; the wrapper renders the user's memo with + # the state-bound prop. + assert f"jsx({definition.export_name}," in page_output + memo_files, _memo_imports = compile_memo_components( + memos=tuple(ctx.auto_memo_components.values()), + ) + wrapper_code = next( + code for path, code in memo_files if definition.export_name in path + ) + assert "useContext(StateContexts" in wrapper_code + assert f"jsx({wrapped.tag}," in wrapper_code + + +def test_user_memo_with_static_props_is_not_auto_memoized() -> None: + """A memo with no reactive props stays inline — no wrapper is generated.""" + + @rx.memo + def static_card(label: rx.Var[str]) -> Component: + return WithProp.create(label=label) + + ctx, page_ctx = _compile_single_page( + lambda: Fragment.create(static_card(label="static")) + ) + + assert not ctx.auto_memo_components + page_output = page_ctx.output_code + assert page_output is not None + assert f"jsx({static_card(label='static').tag}," in page_output + + +def test_user_memo_event_trigger_usecallback_leaves_page_scope() -> None: + """A memo's event-handler prop is memoized inside the generated wrapper. + + An inline arrow recreated on every page render defeats the ``memo`` the + user asked for; the wrapper hoists it into a ``useCallback`` living beside + the state hooks it depends on. + """ + from reflex.compiler.compiler import compile_memo_components + + @rx.memo + def clickable( + on_click: rx.EventHandler[rx.event.no_args_event_spec], + ) -> Component: + return Plain.create(on_click=on_click) + + ctx, page_ctx = _compile_single_page( + lambda: Fragment.create(clickable(on_click=MemoTriggerState.ping)) + ) + + assert not any("useCallback" in hook for hook in page_ctx.hooks) + (definition,) = ctx.auto_memo_components.values() + memo_files, _memo_imports = compile_memo_components( + memos=tuple(ctx.auto_memo_components.values()), + ) + wrapper_code = next( + code for path, code in memo_files if definition.export_name in path + ) + assert "useCallback" in wrapper_code + + +def test_user_memo_children_render_in_page_scope() -> None: + """The wrapper passes children through instead of capturing them. + + Children keep compiling in the page module (so their own reactive parts + get independent wrappers), and the memo body only holds the ``{children}`` + hole plus the state-bound props. + """ + from reflex.compiler.compiler import compile_memo_components + + @rx.memo + def slot_card(label: rx.Var[str], children: rx.Var[Component]) -> Component: + return WithProp.create(children, label=label) + + ctx, page_ctx = _compile_single_page( + lambda: Fragment.create( + slot_card( + Plain.create(LiteralVar.create("static child")), + label=SpecialFormMemoState.value, + ) + ) + ) + + page_output = page_ctx.output_code + assert page_output is not None + assert "useContext(StateContexts" not in page_output + assert 'jsx(Plain,{},"static child")' in page_output + + wrapper_definition = next( + definition + for definition in ctx.auto_memo_components.values() + if isinstance(definition, MemoComponentDefinition) + and isinstance(definition.component, MemoComponent) + ) + memo_files, _memo_imports = compile_memo_components( + memos=tuple(ctx.auto_memo_components.values()), + ) + wrapper_code = next( + code for path, code in memo_files if wrapper_definition.export_name in path + ) + inner_tag = wrapper_definition.component.tag + assert f"jsx({inner_tag}," in wrapper_code + # The hole, not the authored child, is what the memo body renders. + assert ",children)" in wrapper_code + assert "static child" not in wrapper_code + + +def test_user_memo_inside_foreach_is_not_independently_memoized() -> None: + """Foreach owns its snapshot, so a memo inside it renders in that body.""" + from reflex.compiler.compiler import compile_memo_components + + @rx.memo + def row(label: rx.Var[str]) -> Component: + return WithProp.create(label=label) + + ctx, _page_ctx = _compile_single_page( + lambda: rx.box( + rx.foreach(SpecialFormMemoState.items, lambda item: row(label=item)) + ) + ) + + (definition,) = ctx.auto_memo_components.values() + assert isinstance(definition, MemoComponentDefinition) + assert isinstance(definition.component, Foreach) + memo_files, _memo_imports = compile_memo_components( + memos=tuple(ctx.auto_memo_components.values()), + ) + memo_code = "\n".join(code for _, code in memo_files) + assert f"jsx({row(label='x').tag}," in memo_code + + def test_passthrough_memo_skips_hole_for_childless_component() -> None: """Childless components own their JSX output, so the wrapper must not inject a ``{children}`` hole.