feat: add experimental memo decorator for JS-level component and function memoization (#6192)

* feat: add experimental memo decorator for JS-level component and function memoization

Introduce rx.experimental.memo (rx._x.memo) that allows memoizing
components and plain functions at the JavaScript level. Supports
component memos with typed props (including children and rest props
via RestProp), and function memos that emit raw JS. Updates the
compiler pipeline to handle both memo kinds alongside existing
CustomComponent memos, and refactors signature rendering to use
DestructuredArg.

* fix: prevent memo name collisions and compile-time mutation of stored components

Add registry helpers that detect duplicate exported names across memo
kinds and raise on collision. Deepcopy the component before applying
styles during compilation so the stored definition stays clean. Simplify
the function wrappers .call to alias the wrapper itself.

* test: clear old memos when testing.

* test: cleanup

* pyi: update hashes

* fix: camelCase rest-prop keys in memo function bindings and clean up memo internals

Convert remaining_props keys to camelCase in _bind_function_runtime_args
so rest props (e.g. class_name → className) match the component memo
behavior. Also make MemoParam kw_only, return a tuple from get_props
instead of a dict, and remove unnecessary monkeypatch boilerplate from
the integration test fixture.

* refactor: replace memo wrapper closures with proper callable classes

Replace _create_function_wrapper and _create_component_wrapper closures
with _ExperimentalMemoFunctionWrapper and _ExperimentalMemoComponentWrapper
classes, eliminating object.__setattr__ hacks for call/partial/_as_var
in favor of real methods.

* updated hashes

* fix: accept Var[Component] return from component-returning memos

Extract _normalize_component_return to wrap Var[Component] values in
Bare.create, allowing memos that return rx.cond or other component-typed
vars to be registered as component memos. Add a cond overload for
(Any, Var[Component], Var[Component]) -> Component.

* refactor: create per-memo component subclasses with tag set at class level

Replace instance-level self.tag assignment with cached dynamically
created ExperimentalMemoComponent subclasses via
_get_experimental_memo_component_class, so the tag is a class-level
attribute rather than set in _post_init.

---------

Co-authored-by: Masen Furer <m_github@0x26.net>

Author: FarhanAliRaza

pyi_hashes.json

@@ -1,5 +1,5 @@
 {
-  "reflex/__init__.pyi": "a0266c47111e9af7f340186013c7a31e",
+  "reflex/__init__.pyi": "9823934f0e3fca36228004a6fbb1d8df",
   "reflex/components/__init__.pyi": "ac05995852baa81062ba3d18fbc489fb",
   "reflex/components/base/__init__.pyi": "16e47bf19e0d62835a605baa3d039c5a",
   "reflex/components/base/app_wrap.pyi": "22e94feaa9fe675bcae51c412f5b67f1",
@@ -118,5 +118,6 @@
   "reflex/components/recharts/general.pyi": "9abf71810a5405fd45b13804c0a7fd1a",
   "reflex/components/recharts/polar.pyi": "ea4743e8903365ba95bc4b653c47cc4a",
   "reflex/components/recharts/recharts.pyi": "b3d93d085d51053bbb8f65326f34a299",
-  "reflex/components/sonner/toast.pyi": "636050fcc919f8ab0903c30dceaa18f1"
+  "reflex/components/sonner/toast.pyi": "636050fcc919f8ab0903c30dceaa18f1",
+  "reflex/experimental/memo.pyi": "78b1968972194785f72eab32476bc61d"
 }

reflex/__init__.py

@@ -347,7 +347,7 @@
     "utils.imports": ["ImportDict", "ImportVar"],
     "utils.misc": ["run_in_thread"],
     "utils.serializers": ["serializer"],
-    "vars": ["Var", "field", "Field"],
+    "vars": ["Var", "field", "Field", "RestProp"],
 }
 
 _SUBMODULES: set[str] = {

reflex/app.py

@@ -82,6 +82,7 @@
     get_hydrate_event,
     noop,
 )
+from reflex.experimental.memo import EXPERIMENTAL_MEMOS
 from reflex.istate.manager import StateModificationContext
 from reflex.istate.proxy import StateProxy
 from reflex.page import DECORATED_PAGES
@@ -1280,7 +1281,10 @@ def memoized_toast_provider():
             memo_components_output,
             memo_components_result,
             memo_components_imports,
-        ) = compiler.compile_memo_components(dict.fromkeys(CUSTOM_COMPONENTS.values()))
+        ) = compiler.compile_memo_components(
+            dict.fromkeys(CUSTOM_COMPONENTS.values()),
+            tuple(EXPERIMENTAL_MEMOS.values()),
+        )
         compile_results.append((memo_components_output, memo_components_result))
         all_imports.update(memo_components_imports)
         progress.advance(task)

reflex/compiler/compiler.py

@@ -22,6 +22,11 @@
 from reflex.constants.compiler import PageNames, ResetStylesheet
 from reflex.constants.state import FIELD_MARKER
 from reflex.environment import environment
+from reflex.experimental.memo import (
+    ExperimentalMemoComponentDefinition,
+    ExperimentalMemoDefinition,
+    ExperimentalMemoFunctionDefinition,
+)
 from reflex.state import BaseState
 from reflex.style import SYSTEM_COLOR_MODE
 from reflex.utils import console, path_ops
@@ -339,28 +344,46 @@ def _compile_component(component: Component | StatefulComponent) -> str:
 
 def _compile_memo_components(
     components: Iterable[CustomComponent],
+    experimental_memos: Iterable[ExperimentalMemoDefinition] = (),
 ) -> tuple[str, dict[str, list[ImportVar]]]:
     """Compile the components.
 
     Args:
         components: The components to compile.
+        experimental_memos: The experimental memos to compile.
 
     Returns:
         The compiled components.
     """
-    imports = {
-        "react": [ImportVar(tag="memo")],
-        f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")],
-    }
+    imports: dict[str, list[ImportVar]] = {}
     component_renders = []
+    function_renders = []
 
     # Compile each component.
     for component in components:
         component_render, component_imports = utils.compile_custom_component(component)
         component_renders.append(component_render)
         imports = utils.merge_imports(imports, component_imports)
 
-    _apply_common_imports(imports)
+    for memo in experimental_memos:
+        if isinstance(memo, ExperimentalMemoComponentDefinition):
+            memo_render, memo_imports = utils.compile_experimental_component_memo(memo)
+            component_renders.append(memo_render)
+            imports = utils.merge_imports(imports, memo_imports)
+        elif isinstance(memo, ExperimentalMemoFunctionDefinition):
+            memo_render, memo_imports = utils.compile_experimental_function_memo(memo)
+            function_renders.append(memo_render)
+            imports = utils.merge_imports(imports, memo_imports)
+
+    if component_renders:
+        imports = utils.merge_imports(
+            {
+                "react": [ImportVar(tag="memo")],
+                f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")],
+            },
+            imports,
+        )
+        _apply_common_imports(imports)
 
     dynamic_imports = {
         comp_import: None
@@ -380,6 +403,7 @@ def _compile_memo_components(
         templates.memo_components_template(
             imports=utils.compile_imports(imports),
             components=component_renders,
+            functions=function_renders,
             dynamic_imports=sorted(dynamic_imports),
             custom_codes=custom_codes,
         ),
@@ -573,11 +597,13 @@ def compile_page(path: str, component: BaseComponent) -> tuple[str, str]:
 
 def compile_memo_components(
     components: Iterable[CustomComponent],
+    experimental_memos: Iterable[ExperimentalMemoDefinition] = (),
 ) -> tuple[str, str, dict[str, list[ImportVar]]]:
     """Compile the custom components.
 
     Args:
         components: The custom components to compile.
+        experimental_memos: The experimental memos to compile.
 
     Returns:
         The path and code of the compiled components.
@@ -586,7 +612,7 @@ def compile_memo_components(
     output_path = utils.get_components_path()
 
     # Compile the components.
-    code, imports = _compile_memo_components(components)
+    code, imports = _compile_memo_components(components, experimental_memos)
     return output_path, code, imports
 
 

reflex/compiler/templates.py

@@ -8,7 +8,6 @@
 
 from reflex import constants
 from reflex.constants import Hooks
-from reflex.constants.state import CAMEL_CASE_MEMO_MARKER
 from reflex.utils.format import format_state_name, json_dumps
 from reflex.vars.base import VarData
 
@@ -661,6 +660,7 @@ def stateful_components_template(imports: list[_ImportDict], memoized_code: str)
 def memo_components_template(
     imports: list[_ImportDict],
     components: list[dict[str, Any]],
+    functions: list[dict[str, Any]],
     dynamic_imports: Iterable[str],
     custom_codes: Iterable[str],
 ) -> str:
@@ -669,6 +669,7 @@ def memo_components_template(
     Args:
         imports: List of import statements.
         components: List of component definitions.
+        functions: List of function definitions.
         dynamic_imports: List of dynamic import statements.
         custom_codes: List of custom code snippets.
 
@@ -682,21 +683,29 @@ def memo_components_template(
     components_code = ""
     for component in components:
         components_code += f"""
-export const {component["name"]} = memo(({{ {",".join([f"{prop}:{prop}{CAMEL_CASE_MEMO_MARKER}" for prop in component.get("props", [])])} }}) => {{
+export const {component["name"]} = memo(({component["signature"]}) => {{
     {_render_hooks(component.get("hooks", {}))}
     return(
         {_RenderUtils.render(component["render"])}
     )
 }});
 """
 
+    functions_code = ""
+    for function in functions:
+        functions_code += (
+            f"\nexport const {function['name']} = {function['function']};\n"
+        )
+
     return f"""
 {imports_str}
 
 {dynamic_imports_str}
 
 {custom_code_str}
 
+{functions_code}
+
 {components_code}"""
 
 

reflex/compiler/utils.py

@@ -4,6 +4,7 @@
 
 import asyncio
 import concurrent.futures
+import copy
 import operator
 import traceback
 from collections.abc import Mapping, Sequence
@@ -20,14 +21,19 @@
 from reflex.components.el.elements.metadata import Head, Link, Meta, Title
 from reflex.components.el.elements.other import Html
 from reflex.components.el.elements.sectioning import Body
-from reflex.constants.state import FIELD_MARKER
+from reflex.constants.state import CAMEL_CASE_MEMO_MARKER, FIELD_MARKER
+from reflex.experimental.memo import (
+    ExperimentalMemoComponentDefinition,
+    ExperimentalMemoFunctionDefinition,
+)
 from reflex.istate.storage import Cookie, LocalStorage, SessionStorage
 from reflex.state import BaseState, _resolve_delta
 from reflex.style import Style
 from reflex.utils import format, imports, path_ops
 from reflex.utils.imports import ImportVar, ParsedImportDict
 from reflex.utils.prerequisites import get_web_dir
 from reflex.vars.base import Field, Var, VarData
+from reflex.vars.function import DestructuredArg
 
 # To re-export this function.
 merge_imports = imports.merge_imports
@@ -344,6 +350,9 @@ def compile_custom_component(
         {
             "name": component.tag,
             "props": props,
+            "signature": DestructuredArg(
+                fields=tuple(f"{prop}:{prop}{CAMEL_CASE_MEMO_MARKER}" for prop in props)
+            ).to_javascript(),
             "render": render.render(),
             "hooks": render._get_all_hooks(),
             "custom_code": render._get_all_custom_code(),
@@ -353,6 +362,104 @@ def compile_custom_component(
     )
 
 
+def _apply_component_style_for_compile(component: Component) -> Component:
+    """Apply the app style to a compiled component tree.
+
+    Args:
+        component: The component tree.
+
+    Returns:
+        The styled component tree.
+    """
+    try:
+        from reflex.utils.prerequisites import get_and_validate_app
+
+        style = get_and_validate_app().app.style
+    except Exception:
+        style = {}
+
+    component._add_style_recursive(style)
+    return component
+
+
+def compile_experimental_component_memo(
+    definition: ExperimentalMemoComponentDefinition,
+) -> tuple[dict, ParsedImportDict]:
+    """Compile an experimental memo component.
+
+    Args:
+        definition: The component memo definition.
+
+    Returns:
+        A tuple of the compiled component definition and its imports.
+    """
+    render = _apply_component_style_for_compile(copy.deepcopy(definition.component))
+
+    imports: ParsedImportDict = {
+        lib: fields
+        for lib, fields in render._get_all_imports().items()
+        if lib != f"$/{constants.Dirs.COMPONENTS_PATH}"
+    }
+
+    imports.setdefault("@emotion/react", []).append(ImportVar("jsx"))
+
+    signature_fields = [
+        f"{param.js_prop_name}:{param.placeholder_name}"
+        for param in definition.params
+        if not param.is_children and not param.is_rest
+    ]
+
+    if any(param.is_children for param in definition.params):
+        signature_fields.insert(0, "children")
+
+    rest_param = next((param for param in definition.params if param.is_rest), None)
+
+    return (
+        {
+            "kind": "component",
+            "name": definition.export_name,
+            "signature": DestructuredArg(
+                fields=tuple(signature_fields),
+                rest=rest_param.placeholder_name if rest_param is not None else None,
+            ).to_javascript(),
+            "render": render.render(),
+            "hooks": render._get_all_hooks(),
+            "custom_code": render._get_all_custom_code(),
+            "dynamic_imports": render._get_all_dynamic_imports(),
+        },
+        imports,
+    )
+
+
+def compile_experimental_function_memo(
+    definition: ExperimentalMemoFunctionDefinition,
+) -> tuple[dict, ParsedImportDict]:
+    """Compile an experimental memo function.
+
+    Args:
+        definition: The function memo definition.
+
+    Returns:
+        A tuple of the compiled function definition and its imports.
+    """
+    imports: ParsedImportDict = {}
+    if var_data := definition.function._get_all_var_data():
+        imports = {
+            lib: list(fields)
+            for lib, fields in dict(var_data.imports).items()
+            if lib != f"$/{constants.Dirs.COMPONENTS_PATH}"
+        }
+
+    return (
+        {
+            "kind": "function",
+            "name": definition.python_name,
+            "function": str(definition.function),
+        },
+        imports,
+    )
+
+
 def create_document_root(
     head_components: Sequence[Component] | None = None,
     html_lang: str | None = None,

reflex/components/core/cond.py

@@ -98,6 +98,10 @@ def cond(condition: Any, c1: Component, c2: Any, /) -> Component: ...  # pyright
 def cond(condition: Any, c1: Component, /) -> Component: ...
 
 
+@overload
+def cond(condition: Any, c1: Var[Component], c2: Var[Component], /) -> Component: ...  # pyright: ignore [reportOverlappingOverload]
+
+
 @overload
 def cond(condition: Any, c1: Any, c2: Component, /) -> Component: ...  # pyright: ignore [reportOverlappingOverload]
 

reflex/experimental/__init__.py

@@ -8,6 +8,7 @@
 
 from . import hooks as hooks
 from .client_state import ClientStateVar as ClientStateVar
+from .memo import memo as memo
 
 
 class ExperimentalNamespace(SimpleNamespace):
@@ -58,4 +59,5 @@ def register_component_warning(component_name: str):
     client_state=ClientStateVar.create,
     hooks=hooks,
     code_block=code_block,
+    memo=memo,
 )