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

pyi_hashes.json

@@ -1,5 +1,5 @@
 {
-  "reflex/__init__.pyi": "0a3ae880e256b9fd3b960e12a2cb51a7",
+  "reflex/__init__.pyi": "224964f24351c79614ab9a4ae47560a0",
   "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": "d87ff9b85b2a204be01753690df4fb11",
   "reflex/components/recharts/polar.pyi": "b8b1a3e996e066facdf4f8c9eb363137",
   "reflex/components/recharts/recharts.pyi": "d5c9fc57a03b419748f0408c23319eee",
-  "reflex/components/sonner/toast.pyi": "dca44901640cda9d58c62ff8434faa3e"
+  "reflex/components/sonner/toast.pyi": "dca44901640cda9d58c62ff8434faa3e",
+  "reflex/experimental/memo.pyi": "a1c5c4682fc4dadbd82a0a5e8fd4bd32"
 }

reflex/__init__.py

@@ -342,7 +342,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

@@ -84,6 +84,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
@@ -1321,7 +1322,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/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,
 )