From 8d47b987d72a5690fbc343d0efe9a4cf12d324e7 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 11 Jul 2026 02:48:25 +0000 Subject: [PATCH 1/2] Fuse snapshot memo subtree walks into one traversal Compiling a snapshot memo definition walked the (deep-copied) subtree four times - _get_all_hooks, _get_all_custom_code, _get_all_dynamic_imports, _get_all_imports - each independently recursing children and prop components and re-calling per-node getters like _get_components_in_props. _collect_subtree_artifacts gathers all four in a single traversal, combining each artifact at every node in the same order as its dedicated recursion (hooks stay structural-tree-only; prop-subtree custom code still lands before the node's own add_custom_code). The walk runs before render() so add_hooks side effects land first, same as the legacy hooks-walk-then-render order. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Jy8uHH11KircGa2MbTrR8g --- reflex/compiler/utils.py | 84 +++++++++++++++- tests/units/compiler/test_compiler_utils.py | 103 +++++++++++++++++++- 2 files changed, 182 insertions(+), 5 deletions(-) diff --git a/reflex/compiler/utils.py b/reflex/compiler/utils.py index fb537e4d176..3b5339d3a9d 100644 --- a/reflex/compiler/utils.py +++ b/reflex/compiler/utils.py @@ -377,6 +377,80 @@ def _app_style() -> ComponentStyle | Style: return {} +def _collect_subtree_artifacts( + component: Component, + in_prop_tree: bool = False, +) -> tuple[dict[str, VarData | None], dict[str, None], set[str], ParsedImportDict]: + """Collect hooks, custom code, dynamic imports, and imports in one walk. + + Equivalent to calling ``_get_all_hooks``, ``_get_all_custom_code``, + ``_get_all_dynamic_imports``, and ``_get_all_imports`` on ``component``, + but traverses the subtree once: each artifact is combined at every node + in the same order as its dedicated recursion, while per-node getters + (notably ``_get_components_in_props``) run once per node instead of once + per artifact walk. + + Args: + component: The root of the subtree to collect from. + in_prop_tree: Whether ``component`` was reached through a prop rather + than structural children. Hooks are only collected from the + structural tree, matching ``_get_all_hooks``. + + Returns: + A tuple of (hooks, custom code, dynamic imports, imports). + """ + hooks: dict[str, VarData | None] = {} + if not in_prop_tree: + hooks.update(component._get_hooks_internal()) + component_hooks = component._get_hooks() + if component_hooks is not None: + hooks[component_hooks] = None + hooks.update(component._get_added_hooks()) + + custom_code: dict[str, None] = {} + component_custom_code = component._get_custom_code() + if component_custom_code is not None: + custom_code[component_custom_code] = None + + dynamic_imports: set[str] = set() + dynamic_import = component._get_dynamic_imports() + if dynamic_import: + dynamic_imports.add(dynamic_import) + + prop_parts = [ + _collect_subtree_artifacts(prop_component, in_prop_tree=True) + for prop_component in component._get_components_in_props() + ] + # Custom code pulls prop subtrees in before the component's own + # ``add_custom_code`` contributions, matching ``_get_all_custom_code``. + for _, prop_custom_code, _, _ in prop_parts: + custom_code |= prop_custom_code + + for clz in component._iter_parent_classes_with_method("add_custom_code"): + for item in clz.add_custom_code(component): + custom_code[item] = None + + child_parts = [ + _collect_subtree_artifacts(child, in_prop_tree=in_prop_tree) + for child in component.children + ] + for child_hooks, child_custom_code, child_dynamic_imports, _ in child_parts: + hooks.update(child_hooks) + custom_code |= child_custom_code + dynamic_imports |= child_dynamic_imports + + for _, _, prop_dynamic_imports, _ in prop_parts: + dynamic_imports |= prop_dynamic_imports + + all_imports = imports.merge_parsed_imports( + component._get_imports(), + *(child_imports for _, _, _, child_imports in child_parts), + *(prop_imports for _, _, _, prop_imports in prop_parts), + ) + + return hooks, custom_code, dynamic_imports, all_imports + + def compile_experimental_component_memo( definition: MemoComponentDefinition, ) -> tuple[dict, ParsedImportDict]: @@ -417,11 +491,13 @@ def compile_experimental_component_memo( rendered = render.render() else: render = _apply_component_style_for_compile(copy.deepcopy(definition.component)) - hooks = render._get_all_hooks() + # One fused walk replaces the four independent subtree recursions. + # It runs before ``render()`` so ``add_hooks`` side effects (derived + # props) land first, matching the legacy hooks-walk-then-render order. + hooks, custom_code, dynamic_imports, all_imports = _collect_subtree_artifacts( + render + ) rendered = render.render() - custom_code = render._get_all_custom_code() - dynamic_imports = render._get_all_dynamic_imports() - all_imports = render._get_all_imports() # Each un-mirrored memo lives in ``web/utils/components/.jsx`` and is # imported from ``$/utils/components/``. Strip a self-import so a memo diff --git a/tests/units/compiler/test_compiler_utils.py b/tests/units/compiler/test_compiler_utils.py index db7eff1a495..70dc1e56ade 100644 --- a/tests/units/compiler/test_compiler_utils.py +++ b/tests/units/compiler/test_compiler_utils.py @@ -3,8 +3,10 @@ import asyncio import pytest +from reflex_base.components.component import Component +from reflex_base.components.component import field as component_field -from reflex.compiler.utils import compile_state +from reflex.compiler.utils import _collect_subtree_artifacts, compile_state from reflex.constants.state import FIELD_MARKER from reflex.state import State from reflex.vars.base import computed_var @@ -48,3 +50,102 @@ async def test_compile_state_resolves_async_computed_vars_with_running_event_loo assert values[f"a{FIELD_MARKER}"] == 1 assert values[f"b{FIELD_MARKER}"] == 2 assert values[f"async_value{FIELD_MARKER}"] == "resolved" + + +class ArtifactChild(Component): + """Leaf component contributing hooks, custom code, and a dynamic import.""" + + library = "artifact-child-lib" + tag = "ArtifactChild" + + def add_hooks(self) -> list[str]: + """Contribute a hook line. + + Returns: + The hook lines for this component. + """ + return ["const artifactChildHook = 1;"] + + def _get_custom_code(self) -> str | None: + return "const artifactChildCustom = 1;" + + def _get_dynamic_imports(self) -> str | None: + return "artifact-child-dynamic" + + +class ArtifactProp(Component): + """Component used inside a prop slot, with its own custom code.""" + + library = "artifact-prop-lib" + tag = "ArtifactProp" + + def add_hooks(self) -> list[str]: + """Contribute a hook line (must not leak out of the prop tree). + + Returns: + The hook lines for this component. + """ + return ["const artifactPropHook = 1;"] + + def _get_custom_code(self) -> str | None: + return "const artifactPropCustom = 1;" + + +class ArtifactRoot(Component): + """Root component with a component-typed prop slot and class custom code.""" + + library = "artifact-root-lib" + tag = "ArtifactRoot" + + slot: Component | None = component_field(default=None) + + def add_hooks(self) -> list[str]: + """Contribute a hook line. + + Returns: + The hook lines for this component. + """ + return ["const artifactRootHook = 1;"] + + def add_custom_code(self) -> list[str]: + """Contribute class-level custom code. + + Returns: + The custom code lines for this component. + """ + return ["const artifactRootAddCustom = 1;"] + + +def test_collect_subtree_artifacts_matches_legacy_walks(): + """The fused artifact walk matches the four legacy recursions exactly. + + Order-sensitive: hooks and custom code dict ordering determines the + emitted JS, so the fused walk must reproduce each legacy recursion's + insertion order (including custom code from prop subtrees landing before + the component's own ``add_custom_code`` contributions, and hooks never + being collected from prop subtrees). + """ + root = ArtifactRoot.create( + ArtifactChild.create(), + ArtifactChild.create(id="artifact-ref"), + slot=ArtifactProp.create(ArtifactChild.create()), + ) + + fused_hooks, fused_custom, fused_dynamic, fused_imports = ( + _collect_subtree_artifacts(root) + ) + + legacy_hooks = root._get_all_hooks() + assert fused_hooks == legacy_hooks + assert list(fused_hooks) == list(legacy_hooks) + assert "const artifactPropHook = 1;" not in fused_hooks + + legacy_custom = root._get_all_custom_code() + assert fused_custom == legacy_custom + assert list(fused_custom) == list(legacy_custom) + + assert fused_dynamic == root._get_all_dynamic_imports() + + legacy_imports = root._get_all_imports() + assert fused_imports == legacy_imports + assert list(fused_imports) == list(legacy_imports) From 82b57dd0f70c4e128ae45a62248eaace8425f6bc Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 11 Jul 2026 02:49:03 +0000 Subject: [PATCH 2/2] Add news fragment for snapshot memo walk fusion Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Jy8uHH11KircGa2MbTrR8g --- news/6748.performance.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 news/6748.performance.md diff --git a/news/6748.performance.md b/news/6748.performance.md new file mode 100644 index 00000000000..e009d09df71 --- /dev/null +++ b/news/6748.performance.md @@ -0,0 +1 @@ +Snapshot memo definitions collect hooks, custom code, dynamic imports, and imports in one subtree traversal instead of four independent recursive walks.