374 lines
15 KiB
Python
374 lines
15 KiB
Python
"""MemoizeStatefulPlugin — auto-memoize stateful components with ``rx._x.memo``.
|
|
|
|
This plugin replaces the legacy ``StatefulComponent`` wrapping pass. It
|
|
participates in the normal single-pass walk via ``enter_component`` and inserts
|
|
per-subtree ``{children}``-pass-through wrappers built on the experimental
|
|
memo infrastructure. The wrapped subtree remains in the tree for the normal
|
|
walker descent, so downstream plugins (e.g. ``DefaultCollectorPlugin``) still
|
|
see the original components and collect their imports/hooks as usual.
|
|
|
|
Each unique subtree shape contributes:
|
|
|
|
- One generated experimental memo component definition, compiled into the
|
|
shared ``$/utils/components`` module.
|
|
- ``useCallback`` hook lines for each non-lifecycle event trigger, emitted into
|
|
the generated memo body so handler hooks stay inside that rendering domain.
|
|
|
|
No shared ``stateful_components`` file is produced.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
from typing import Any
|
|
|
|
from reflex_base.components.component import BaseComponent, Component
|
|
from reflex_base.components.memoize_helpers import (
|
|
MemoizationStrategy,
|
|
_is_structural_memoization_child,
|
|
fix_event_triggers_for_memo,
|
|
get_memoization_strategy,
|
|
is_snapshot_boundary,
|
|
)
|
|
from reflex_base.constants.compiler import MemoizationDisposition
|
|
from reflex_base.plugins import ComponentAndChildren, PageContext
|
|
from reflex_base.plugins.base import Plugin
|
|
|
|
from reflex.experimental.memo import create_passthrough_component_memo
|
|
|
|
|
|
def _subtree_has_reactive_data(
|
|
component: Component, _cache: dict[int, bool] | None = None
|
|
) -> bool:
|
|
"""Whether ``component``'s subtree carries reactive signals worth memoizing.
|
|
|
|
No-arg event handlers (``on_click=State.ping``) contribute hooks only via
|
|
``event_triggers`` / ``_get_events_hooks``, not as a Var, so the per-Var
|
|
scan must be paired with an explicit ``event_triggers`` check.
|
|
|
|
``useRef`` from a static ``id`` prop is intentionally ignored — it lives
|
|
in ``_get_hooks_internal``, not in any Var, so static-id-only elements
|
|
don't surface here and aren't flagged.
|
|
|
|
Args:
|
|
component: The component whose subtree to inspect.
|
|
_cache: Internal ``id()``-keyed cache of per-subtree results so
|
|
components reachable via overlapping ``var_data.components`` and
|
|
``children`` paths are evaluated once. ``False`` is also used as
|
|
a transient placeholder while a subtree is being computed to
|
|
break cycles.
|
|
|
|
Returns:
|
|
True if the subtree carries event triggers, explicit hooks, or any
|
|
Var whose merged var_data has ``state`` or ``hooks``.
|
|
"""
|
|
if _cache is None:
|
|
_cache = {}
|
|
key = id(component)
|
|
cached = _cache.get(key)
|
|
if cached is not None:
|
|
return cached
|
|
# Placeholder breaks cycles: a subtree that references itself is
|
|
# treated as non-reactive on the recursive arm; the real result for
|
|
# this node is written back below.
|
|
_cache[key] = False
|
|
result = _component_subtree_is_reactive(component, _cache)
|
|
_cache[key] = result
|
|
return result
|
|
|
|
|
|
def _component_subtree_is_reactive(
|
|
component: Component, _cache: dict[int, bool]
|
|
) -> bool:
|
|
"""Inner walk for :func:`_subtree_has_reactive_data` (uncached node check).
|
|
|
|
Internal hooks (``_get_hooks_internal``) cover event-trigger callbacks,
|
|
lifecycle hooks (``on_mount``/``on_unmount``), and Var-derived hooks
|
|
(state context, client state, custom). The static ``id`` ref hook is
|
|
explicitly subtracted so an id-only element does not flag as reactive.
|
|
|
|
Args:
|
|
component: The component to inspect.
|
|
_cache: Shared cache passed through recursive calls.
|
|
|
|
Returns:
|
|
True if ``component`` itself or any reachable descendant carries
|
|
reactive signals.
|
|
"""
|
|
ref_hook = component._get_ref_hook()
|
|
ref_hook_key = str(ref_hook) if ref_hook is not None else None
|
|
for hook_key in component._get_hooks_internal():
|
|
if hook_key != ref_hook_key:
|
|
return True
|
|
if component._get_hooks() is not None or component._get_added_hooks():
|
|
return True
|
|
for var in component._get_vars(include_children=False):
|
|
var_data = var._get_all_var_data()
|
|
if var_data is None:
|
|
continue
|
|
if var_data.state or var_data.hooks:
|
|
return True
|
|
for comp in var_data.components:
|
|
if isinstance(comp, Component) and _subtree_has_reactive_data(comp, _cache):
|
|
return True
|
|
for child in component.children:
|
|
if isinstance(child, Component) and _subtree_has_reactive_data(child, _cache):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _should_memoize(component: Component) -> bool:
|
|
"""Decide whether ``component`` is a candidate for auto-memoization.
|
|
|
|
Snapshot boundaries (``recursive=False``) suppress their descendants,
|
|
so a stateful subtree must trigger wrapping at the boundary itself —
|
|
otherwise the state read leaks into the page module. Other components
|
|
are evaluated from their own props/triggers; descendants are visited
|
|
independently by the walker.
|
|
|
|
Args:
|
|
component: The candidate component.
|
|
|
|
Returns:
|
|
True if the component should be wrapped in a memo definition.
|
|
"""
|
|
from reflex_components_core.base.bare import Bare
|
|
from reflex_components_core.core.cond import Cond
|
|
from reflex_components_core.core.match import Match
|
|
|
|
strategy = get_memoization_strategy(component)
|
|
|
|
if component._memoization_mode.disposition == MemoizationDisposition.NEVER:
|
|
return False
|
|
if isinstance(component, Bare):
|
|
# A stateful value will be wrapped in a separate component. Match the
|
|
# per-Var predicate used by ``_subtree_has_reactive_data`` so a Bare
|
|
# whose Var carries only imports (no state/hooks) is not memoized.
|
|
contents_var_data = component.contents._get_all_var_data()
|
|
if contents_var_data is not None:
|
|
if contents_var_data.state or contents_var_data.hooks:
|
|
return True
|
|
for embedded in contents_var_data.components:
|
|
if isinstance(embedded, Component) and _subtree_has_reactive_data(
|
|
embedded
|
|
):
|
|
return True
|
|
# Cond and Match render conditional branch JSX from their own props rather
|
|
# than from a tag; structural memoization children (e.g. ``Foreach``)
|
|
# render via their own structural form. All have no ``tag`` but still
|
|
# must be considered for memoization — the structural-child case in
|
|
# particular owns its whole subtree as a snapshot, so if it does not
|
|
# wrap here, its descendants leak to the page module.
|
|
if (
|
|
component.tag is None
|
|
and not isinstance(component, (Cond, Match))
|
|
and not _is_structural_memoization_child(component)
|
|
):
|
|
return False
|
|
if component._memoization_mode.disposition == MemoizationDisposition.ALWAYS:
|
|
return True
|
|
|
|
# Direct Vars only (component's own props, style, class_name, id, etc.).
|
|
# Match the per-Var predicate used by ``_subtree_has_reactive_data``
|
|
# var_data carrying only imports is not reactive.
|
|
for prop_var in component._get_vars(include_children=False):
|
|
var_data = prop_var._get_all_var_data()
|
|
if var_data is None:
|
|
continue
|
|
if var_data.state or var_data.hooks:
|
|
return True
|
|
for embedded in var_data.components:
|
|
if isinstance(embedded, Component) and _subtree_has_reactive_data(embedded):
|
|
return True
|
|
|
|
if strategy is MemoizationStrategy.SNAPSHOT and not is_snapshot_boundary(component):
|
|
return True
|
|
|
|
if is_snapshot_boundary(component) and _subtree_has_reactive_data(component):
|
|
return True
|
|
|
|
# Components with event triggers are always memoized (to wrap callbacks).
|
|
return bool(component.event_triggers)
|
|
|
|
|
|
@dataclasses.dataclass(frozen=True, slots=True)
|
|
class MemoizeStatefulPlugin(Plugin):
|
|
"""Auto-memoize stateful components with experimental-memo wrappers.
|
|
|
|
Registered in ``default_page_plugins`` before ``DefaultCollectorPlugin``.
|
|
Components either render as passthrough memo wrappers or snapshot memo
|
|
wrappers (see ``get_memoization_strategy``):
|
|
|
|
- Snapshot wrappers (``MemoizationLeaf``-style boundaries and structural
|
|
``Foreach`` wrappers): wrapped in ``enter_component``
|
|
and returned with empty structural children. The walker skips descent, so
|
|
hooks attached to the captured body are compiled into the memo body only.
|
|
- Passthrough wrappers are wrapped in
|
|
``leave_component`` after descendants have already compiled, so any inner
|
|
memo wrappers flow into this wrapper's children.
|
|
|
|
Descendants of a snapshot boundary are never independently memoized; the
|
|
boundary owns the wrapping decision for its whole subtree. This is tracked
|
|
via ``PageContext.memoize_suppressor_stack`` — a stack of component ids
|
|
that pushed suppression, popped in ``leave_component`` when the matching
|
|
component leaves.
|
|
"""
|
|
|
|
def enter_component(
|
|
self,
|
|
comp: BaseComponent,
|
|
/,
|
|
*,
|
|
page_context: PageContext,
|
|
compile_context: Any,
|
|
in_prop_tree: bool = False,
|
|
) -> BaseComponent | ComponentAndChildren | None:
|
|
"""Memoize snapshot-boundary subtrees before descent.
|
|
|
|
Snapshot boundaries (``MemoizationLeaf``-style, see
|
|
``is_snapshot_boundary``) stash state-referencing hooks inside
|
|
internally-built structural children. If we waited until
|
|
``leave_component`` to swap the boundary for its memo wrapper, the
|
|
walker would have already descended and the collector plugin would
|
|
have pulled those hooks into page scope. Returning the wrapper with
|
|
empty structural children here causes the walker to skip the descent
|
|
entirely — the boundary's full snapshot lives only in the memo
|
|
component definition compiled separately.
|
|
|
|
Non-boundary components are handled in ``leave_component`` so their
|
|
already-compiled children flow into the wrapper.
|
|
|
|
Args:
|
|
comp: The component being visited.
|
|
page_context: The active page context.
|
|
compile_context: The active compile context.
|
|
in_prop_tree: Whether the component is in a prop subtree.
|
|
|
|
Returns:
|
|
A ``(wrapper, ())`` replacement for memoized boundaries, otherwise
|
|
``None``.
|
|
"""
|
|
if in_prop_tree:
|
|
return None
|
|
if not isinstance(comp, Component):
|
|
return None
|
|
if page_context.memoize_suppressor_stack:
|
|
return None
|
|
strategy = get_memoization_strategy(comp)
|
|
if strategy is not MemoizationStrategy.SNAPSHOT:
|
|
return None
|
|
snapshot_boundary = is_snapshot_boundary(comp)
|
|
|
|
if not _should_memoize(comp):
|
|
# Boundary not worth wrapping — still suppress descendants so
|
|
# they don't memoize independently of the boundary's subtree.
|
|
if snapshot_boundary:
|
|
page_context.memoize_suppressor_stack.append(id(comp))
|
|
return None
|
|
|
|
wrapper = self._build_wrapper(
|
|
comp,
|
|
page_context,
|
|
compile_context,
|
|
)
|
|
return None if wrapper is None else (wrapper, ())
|
|
|
|
def leave_component(
|
|
self,
|
|
comp: BaseComponent,
|
|
children: tuple[BaseComponent, ...],
|
|
/,
|
|
*,
|
|
page_context: PageContext,
|
|
compile_context: Any,
|
|
in_prop_tree: bool = False,
|
|
) -> BaseComponent | ComponentAndChildren | None:
|
|
"""Wrap non-boundary memoizables and pop any suppression this component pushed.
|
|
|
|
Args:
|
|
comp: The component being visited.
|
|
children: Its compiled children (unused; the wrapper reads from
|
|
``comp.children`` which the walker has already updated).
|
|
page_context: The active page context.
|
|
compile_context: The active compile context.
|
|
in_prop_tree: Whether the component is in a prop subtree.
|
|
|
|
Returns:
|
|
The memo wrapper for non-boundary memoizables, else ``None``.
|
|
"""
|
|
if in_prop_tree:
|
|
return None
|
|
if not isinstance(comp, Component):
|
|
return None
|
|
|
|
stack = page_context.memoize_suppressor_stack
|
|
if stack and stack[-1] == id(comp):
|
|
stack.pop()
|
|
|
|
if stack:
|
|
return None
|
|
|
|
if len(children) != len(comp.children) or any(
|
|
compiled_child is not current_child
|
|
for compiled_child, current_child in zip(
|
|
children, comp.children, strict=True
|
|
)
|
|
):
|
|
comp = page_context.own(comp)
|
|
comp.children = list(children)
|
|
|
|
strategy = get_memoization_strategy(comp)
|
|
if strategy is MemoizationStrategy.SNAPSHOT:
|
|
return None
|
|
|
|
if not _should_memoize(comp):
|
|
return None
|
|
|
|
return self._build_wrapper(comp, page_context, compile_context)
|
|
|
|
@staticmethod
|
|
def _build_wrapper(
|
|
comp: Component,
|
|
page_context: PageContext,
|
|
compile_context: Any,
|
|
) -> BaseComponent | None:
|
|
"""Return the memo wrapper component for ``comp``, or ``None`` if untagged.
|
|
|
|
Rewrites ``comp.event_triggers`` on a page-local clone via
|
|
:func:`fix_event_triggers_for_memo` so the memo body renders the
|
|
memoized ``useCallback`` forms, and registers the memo definition on
|
|
``compile_context`` so the memo module compile pass emits it.
|
|
|
|
Args:
|
|
comp: The component being memoized.
|
|
page_context: The active page context.
|
|
compile_context: The active compile context.
|
|
|
|
Returns:
|
|
The wrapper instance, or ``None`` if the component's render is
|
|
empty and has no meaningful tag.
|
|
"""
|
|
comp = fix_event_triggers_for_memo(comp, page_context)
|
|
|
|
# Passthrough memo definitions capture app-specific event/state vars, so
|
|
# 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()
|
|
# The wrapper has no structural children at the page level, but parents
|
|
# walking ``_get_all_refs`` (e.g. ``Form._get_form_refs`` collecting
|
|
# ref_<id> mappings into ``handleSubmit``) need to see refs from the
|
|
# wrapped subtree. Delegate ref collection to the original component
|
|
# so descendants inside the memo body remain reachable for ref lookup.
|
|
object.__setattr__(wrapper, "_get_all_refs", comp._get_all_refs)
|
|
return wrapper
|
|
|
|
|
|
__all__ = ["MemoizeStatefulPlugin"]
|