Skip to content

Single-pass compiler: Replace StatefulComponent auto-memoization with rx._x.memo-based approach #6213

@masenf

Description

@masenf

Summary

Replace the current StatefulComponent auto-memoization mechanism with one that uses rx._x.memo (the experimental memo decorator) to create stateful components that reference {children} instead of recursively recreating the component tree. Implement this as a compiler plugin in the new single-pass architecture.

Depends on ENG-9142 and ENG-9143.

Background: How StatefulComponent Works Today

StatefulComponent in reflex/components/component.py (~line 2375) is a BaseComponent subclass that wraps components which depend on state. The current flow:

  1. StatefulComponent.compile_from(component) walks the tree recursively
  2. For each Component, it calls StatefulComponent.create(component) which:
    • Checks if the component has _memoization_mode enabled
    • Renders the component to a string and hashes it to create a unique tag name (_get_tag_name)
    • If a component with that tag already exists in tag_to_stateful_component, bumps its references counter
    • Otherwise creates a new StatefulComponent wrapper
  3. StatefulComponent._render_stateful_code() renders the wrapper as a memoized React component using stateful_component_template()
  4. In prod mode, shared components (referenced > 1 time) are extracted to a separate stateful_components.js file via _compile_stateful_components() / _get_shared_components_recursive()

Problems with this approach:

  • Re-renders the entire component subtree inside the memoized wrapper — the children are baked into the rendered code, not passed as {children}
  • Hash-based identity is fragile — rendering to string and hashing is expensive and can produce false duplicates or miss actual duplicates
  • Complex shared component tracking — the references counter and rendered_as_shared flag add complexity
  • Separate tree walkStatefulComponent.compile_from() is yet another full traversal of the component tree, done before the main compilation walks

Proposed Approach

New MemoizeStatefulPlugin compiler plugin

Instead of a separate StatefulComponent.compile_from() pass, implement auto-memoization as a CompilerPlugin that runs during the single tree walk:

class MemoizeStatefulPlugin(CompilerPlugin):
    async def compile_component(self, comp, /, **kwargs):
        """Wrap stateful components with rx._x.memo."""
        # Pre-yield: check if this component depends on state
        comp, children = yield
        
        if not isinstance(comp, Component):
            yield
            return
            
        if not comp._memoization_mode.disposition:
            yield
            return
            
        # Check if component actually references state vars
        if not self._has_state_dependency(comp):
            yield
            return
        
        # Wrap with memo, using {children} placeholder instead of
        # re-rendering the entire subtree
        memo_comp = self._create_memo_wrapper(comp)
        yield memo_comp

Use {children} instead of baking the subtree

The key architectural change: instead of rendering the entire subtree into the memoized component's body, the memo wrapper should accept {children} as a prop. This means:

  • The memoized component renders as <MemoComp_{hash}>{children}</MemoComp_{hash}>
  • Children are passed through, not duplicated
  • React's own reconciliation handles re-rendering children efficiently
  • No need for the references / rendered_as_shared tracking

Use rx._x.memo decorator

The existing memo (alias for custom_component) in component.py creates CustomComponent wrappers. The experimental rx._x.memo should be used instead, which provides proper React.memo semantics. The stateful component wrapper should use this decorator to generate the memoized JS output.

Acceptance Criteria

  • New MemoizeStatefulPlugin implements compile_component hook
  • Memoized components use {children} prop instead of baking the subtree
  • The old StatefulComponent class is removed or deprecated
  • _compile_stateful_components() and _get_shared_components_recursive() in compiler.py are removed
  • The stateful_components.js shared file is no longer needed (or simplified)
  • StatefulComponent.compile_from() tree walk is eliminated
  • Existing apps produce functionally equivalent output (same visual behavior, may differ in JS structure)
  • Unit tests verify that stateful components are properly memoized
  • Integration test: an app with shared stateful components across pages works correctly

Key Files

  • reflex/components/component.py:
    • StatefulComponent class (~line 2375) — the class to replace
    • StatefulComponent.compile_from() (~line 2788) — the tree walk to eliminate
    • StatefulComponent.create() (~line 2411) — the memoization logic
    • StatefulComponent._get_tag_name() (~line 2517) — hash-based naming
    • MemoizationMode — controls which components get memoized
  • reflex/compiler/compiler.py:
    • compile_stateful_components() — orchestrates the memoization pass
    • _compile_stateful_components() — extracts shared components
    • _get_shared_components_recursive() — shared component tracking
  • reflex/compiler/templates.pystateful_component_template()

Notes

  • The MemoizationMode class on components should still be respected — components can opt in/out of memoization.
  • This change may affect how StatefulComponent's special _get_all_* methods work (they currently short-circuit when rendered_as_shared is True). With the new approach, this complexity goes away.
  • Consider whether the new memo wrapper needs to handle event trigger hooks (memo_trigger_hooks on the current StatefulComponent).

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementAnything you want improved

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions