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#6192FarhanAliRaza wants to merge 7 commits intoreflex-dev:mainfrom
Conversation
…tion 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.
Merging this PR will improve performance by 3.28%
|
| Benchmark | BASE |
HEAD |
Efficiency | |
|---|---|---|---|---|
| ⚡ | test_compile_stateful[_stateful_page] |
150.9 µs | 146.1 µs | +3.28% |
Comparing FarhanAliRaza:exp-memo (1501541) with main (7ee3026)
Greptile SummaryThis PR introduces Key changes:
Issues found:
Confidence Score: 3/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User as User Code
participant Memo as @rx._x.memo
participant Registry as EXPERIMENTAL_MEMOS
participant Compiler as compiler.py
participant Utils as compiler/utils.py
participant Template as templates.py
User->>Memo: @rx._x.memo on fn (→ Component or → Var)
Memo->>Memo: _analyze_params(fn)
Memo->>Memo: _evaluate_memo_function(fn, params)
alt Component-returning
Memo->>Memo: _create_component_definition()
Memo->>Memo: _lift_rest_props(component)
Memo->>Registry: _register_memo_definition(ExperimentalMemoComponentDefinition)
Memo-->>User: _create_component_wrapper() → callable
else Var-returning
Memo->>Memo: _create_function_definition()
Memo->>Memo: _validate_var_return_expr()
Memo->>Registry: _register_memo_definition(ExperimentalMemoFunctionDefinition)
Memo-->>User: _create_function_wrapper() → callable
end
User->>User: App.compile()
User->>Compiler: compile_memo_components(CUSTOM_COMPONENTS, EXPERIMENTAL_MEMOS)
loop each ExperimentalMemoComponentDefinition
Compiler->>Utils: compile_experimental_component_memo(definition)
Utils->>Utils: copy.deepcopy(definition.component)
Utils->>Utils: _apply_component_style_for_compile()
Utils-->>Compiler: (render_dict, imports)
end
loop each ExperimentalMemoFunctionDefinition
Compiler->>Utils: compile_experimental_function_memo(definition)
Utils-->>Compiler: (function_dict, imports)
end
Compiler->>Template: memo_components_template(components, functions, ...)
Template-->>Compiler: JS file with export const ... = memo(...) and export const ... = ((...) => ...)
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if _is_component_annotation(return_annotation): | ||
| definition = _create_component_definition(fn, return_annotation) | ||
| EXPERIMENTAL_MEMOS[definition.export_name] = definition | ||
| return _create_component_wrapper(definition) | ||
|
|
||
| if _is_var_annotation(return_annotation): | ||
| definition = _create_function_definition(fn, return_annotation) | ||
| EXPERIMENTAL_MEMOS[definition.python_name] = definition | ||
| return _create_function_wrapper(definition) |
There was a problem hiding this comment.
Component memos are stored in EXPERIMENTAL_MEMOS under their title-cased export_name (e.g. "FooBar"), while var-returning function memos are stored under their raw python_name. This means a function named FooBar (key "FooBar") and a component named foo_bar (key "FooBar" via format.to_title_case) silently overwrite each other with no warning or error:
# Both end up at key "FooBar" – second wins silently
@rx._x.memo
def foo_bar(...) -> rx.Component: ... # key = "FooBar"
@rx._x.memo
def FooBar(...) -> rx.Var[str]: ... # key = "FooBar" ← overwrites the componentConsider asserting the key is not already present and raising a descriptive ValueError, or unifying the key strategy (e.g. always use python_name).
| if component_renders: | ||
| imports = utils.merge_imports( | ||
| { | ||
| "react": [ImportVar(tag="memo")], | ||
| f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")], | ||
| }, | ||
| imports, | ||
| ) | ||
|
|
||
| if component_renders: | ||
| _apply_common_imports(imports) |
There was a problem hiding this comment.
if component_renders: guard
The same condition if component_renders: is checked twice in succession. Both branches should be merged into a single block:
| if component_renders: | |
| imports = utils.merge_imports( | |
| { | |
| "react": [ImportVar(tag="memo")], | |
| f"$/{constants.Dirs.STATE_PATH}": [ImportVar(tag="isTrue")], | |
| }, | |
| imports, | |
| ) | |
| if component_renders: | |
| _apply_common_imports(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) |
| @wraps(definition.fn) | ||
| def wrapper(*args: Any, **kwargs: Any) -> Var: | ||
| return imported_var.call( | ||
| *_bind_function_runtime_args(definition, *args, **kwargs) | ||
| ) | ||
|
|
||
| def call(*args: Any, **kwargs: Any) -> Var: | ||
| return imported_var.call( | ||
| *_bind_function_runtime_args(definition, *args, **kwargs) | ||
| ) | ||
|
|
||
| def partial(*args: Any, **kwargs: Any) -> FunctionVar: | ||
| return imported_var.partial( | ||
| *_bind_function_runtime_args(definition, *args, **kwargs) | ||
| ) | ||
|
|
||
| object.__setattr__(wrapper, "call", call) | ||
| object.__setattr__(wrapper, "partial", partial) | ||
| object.__setattr__(wrapper, "_as_var", lambda: imported_var) | ||
| return wrapper |
There was a problem hiding this comment.
call is an exact duplicate of wrapper
The inner call function (lines 673–676) has an identical body to wrapper (lines 668–671). Since object.__setattr__(wrapper, "call", call) attaches it as wrapper.call, calling memo_fn(...) and memo_fn.call(...) will do the exact same thing. The duplication is unnecessary — call could simply alias wrapper:
call = wrapperOr if you want to keep the name separate for clarity, at minimum add a comment explaining the intent.
| def _lift_rest_props(component: Component) -> Component: | ||
| """Convert RestProp children into special props. | ||
|
|
||
| Args: | ||
| component: The component tree to rewrite. | ||
|
|
||
| Returns: | ||
| The rewritten component tree. | ||
| """ | ||
| special_props = list(component.special_props) | ||
| rewritten_children = [] | ||
|
|
||
| for child in component.children: | ||
| if isinstance(child, Bare) and isinstance(child.contents, RestProp): | ||
| special_props.append(child.contents) | ||
| continue | ||
|
|
||
| if isinstance(child, Component): | ||
| child = _lift_rest_props(child) | ||
|
|
||
| rewritten_children.append(child) | ||
|
|
||
| component.children = rewritten_children | ||
| component.special_props = special_props | ||
| return component |
There was a problem hiding this comment.
_lift_rest_props mutates a component tree shared with the frozen definition
_lift_rest_props performs direct in-place mutation of the component tree (component.children = rewritten_children, component.special_props = special_props) and then returns the same object. That mutated component is stored inside the frozen ExperimentalMemoComponentDefinition:
return ExperimentalMemoComponentDefinition(
...
component=_lift_rest_props(component), # ← same object, now mutated
)Later, compile_experimental_component_memo calls _apply_component_style_for_compile(definition.component), which in turn calls component._add_style_recursive(style) on this shared, mutable object. If compile_experimental_component_memo is ever invoked more than once for the same definition (e.g. during hot-reload, testing, or multi-pass compilation), styles will be applied cumulatively to the stored component tree, producing incorrect output.
A safer approach would be to either deep-copy the component before mutating it, or ensure _apply_component_style_for_compile works on a fresh copy each time.
| self.tag = definition.export_name | ||
|
|
||
| props: dict[str, Any] = {} | ||
| for key, value in {**declared_props, **rest_props}.items(): | ||
| camel_cased_key = format.to_camel_case(key) | ||
| literal_value = LiteralVar.create(value) | ||
| props[camel_cased_key] = literal_value | ||
| setattr(self, camel_cased_key, literal_value) | ||
|
|
||
| prop_names = dict.fromkeys(props) | ||
| object.__setattr__(self, "get_props", lambda: prop_names) |
There was a problem hiding this comment.
tag set as an instance attribute instead of the conventional class-level field
Reflex's Component hierarchy conventionally declares tag as a class variable (e.g. tag = "MyComponent"). Setting it as an instance attribute via self.tag = definition.export_name bypasses the class-level mechanism. While Pydantic allows this with mutable models, it could interact unexpectedly with class-level tag resolution, caching, or type-checking logic elsewhere in the framework.
Consider instead making ExperimentalMemoComponent a dynamic subclass with the tag baked in at class creation time, similarly to how CustomComponent derivatives are constructed.
… 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.
| def _bind_function_runtime_args( | ||
| definition: ExperimentalMemoFunctionDefinition, | ||
| *args: Any, | ||
| **kwargs: Any, | ||
| ) -> tuple[Any, ...]: | ||
| """Bind runtime args for a var-returning memo. | ||
|
|
||
| Args: | ||
| definition: The function memo definition. | ||
| *args: Positional arguments. | ||
| **kwargs: Keyword arguments. | ||
|
|
||
| Returns: | ||
| The ordered arguments for the imported FunctionVar. | ||
|
|
||
| Raises: | ||
| TypeError: If the provided arguments are invalid. | ||
| """ | ||
| children_param = _get_children_param(definition.params) | ||
| rest_param = _get_rest_param(definition.params) | ||
| if "children" in kwargs: | ||
| msg = f"`{definition.python_name}` only accepts children positionally." | ||
| raise TypeError(msg) | ||
|
|
||
| if rest_param is not None and rest_param.name in kwargs: | ||
| msg = ( | ||
| f"`{definition.python_name}` captures rest props from extra keyword " | ||
| f"arguments. Do not pass `{rest_param.name}=` directly." | ||
| ) | ||
| raise TypeError(msg) | ||
|
|
||
| if args and children_param is None: | ||
| msg = f"`{definition.python_name}` only accepts keyword props." | ||
| raise TypeError(msg) | ||
|
|
||
| if any(not _is_component_child(child) for child in args): | ||
| msg = ( | ||
| f"`{definition.python_name}` only accepts positional children that are " | ||
| "`rx.Component` or `rx.Var[rx.Component]`." | ||
| ) | ||
| raise TypeError(msg) | ||
|
|
||
| explicit_params = [ | ||
| param | ||
| for param in definition.params | ||
| if not param.is_rest and not param.is_children | ||
| ] | ||
| explicit_values = {} | ||
| remaining_props = kwargs.copy() | ||
| for param in explicit_params: | ||
| if param.name in remaining_props: | ||
| explicit_values[param.name] = remaining_props.pop(param.name) | ||
| elif param.default is not inspect.Parameter.empty: | ||
| explicit_values[param.name] = param.default | ||
| else: | ||
| msg = f"`{definition.python_name}` is missing required prop `{param.name}`." | ||
| raise TypeError(msg) | ||
|
|
||
| if remaining_props and rest_param is None: | ||
| unexpected_prop = next(iter(remaining_props)) | ||
| msg = ( | ||
| f"`{definition.python_name}` does not accept prop `{unexpected_prop}`. " | ||
| "Only declared props may be passed when no `rx.RestProp` is present." | ||
| ) | ||
| raise TypeError(msg) | ||
|
|
||
| if children_param is None and rest_param is None: | ||
| return tuple(explicit_values[param.name] for param in explicit_params) | ||
|
|
||
| children_value: Any | None = None | ||
| if children_param is not None: | ||
| children_value = args[0] if len(args) == 1 else Fragment.create(*args) | ||
|
|
||
| bound_props = {} | ||
| if children_param is not None: | ||
| bound_props[children_param.name] = children_value | ||
| bound_props.update(explicit_values) | ||
| bound_props.update(remaining_props) | ||
| return (bound_props,) |
There was a problem hiding this comment.
_bind_function_runtime_args has several distinct logical phases (input validation, explicit-prop binding, rest-prop collection, return-value construction) but contains no brief comments separating them. Per the project style guidelines, add blank lines and inline comments at each phase boundary so future maintainers can orient quickly.
For example:
# --- validate positional / reserved keyword usage ---
if "children" in kwargs:
...
# --- bind explicit (non-children, non-rest) props ---
explicit_params = [...]
...
# --- collect remaining props into rest ---
if remaining_props and rest_param is None:
...
# --- build the bound argument tuple ---
if children_param is None and rest_param is None:
return ...The same pattern applies to _create_component_wrapper (lines 762–808).
Rule Used: Add blank lines between logical sections of code f... (source)
Learnt From
reflex-dev/flexgen#2170
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
| prop_names = dict.fromkeys(props) | ||
| object.__setattr__(self, "get_props", lambda: prop_names) |
There was a problem hiding this comment.
get_props override returns dict instead of an iterable of prop-name strings
dict.fromkeys(props) returns a dict[str, None] (keys are camelCased prop names, values are all None). Code that iterates get_props() with for k in self.get_props() works correctly because iterating a dict yields its keys, but any call site that expects a plain list or set (e.g. using in with O(1) semantics, calling .append(), or isinstance(..., list) checks) may behave unexpectedly.
Consider using tuple or frozenset to make the intended semantics explicit:
| prop_names = dict.fromkeys(props) | |
| object.__setattr__(self, "get_props", lambda: prop_names) | |
| prop_names = tuple(props) | |
| object.__setattr__(self, "get_props", lambda: prop_names) |
Or, if the ordering or O(1) lookup matters, use a dict.keys() view. At minimum, add a comment clarifying that only the keys are meaningful.
| children_value: Any | None = None | ||
| if children_param is not None: | ||
| children_value = args[0] if len(args) == 1 else Fragment.create(*args) | ||
|
|
||
| bound_props = {} | ||
| if children_param is not None: | ||
| bound_props[children_param.name] = children_value | ||
| bound_props.update(explicit_values) | ||
| bound_props.update(remaining_props) | ||
| return (bound_props,) |
There was a problem hiding this comment.
For component-returning memos, ExperimentalMemoComponent._post_init (line 110) runs format.to_camel_case(key) on every extra kwarg before storing it as a JSX prop. For function-returning memos, _bind_function_runtime_args merges remaining_props straight into bound_props with no conversion:
bound_props.update(remaining_props) # keys stay as Python snake_caseThis means that a function memo invoked with class_name="foo" will receive {class_name: "foo"} in JavaScript, not {className: "foo"}, while the same kwarg passed to a component memo correctly becomes className. A user who switches from a component memo to a function memo, or who passes standard CSS props as rest props, will get silently wrong JS property names.
Consider applying camelCase conversion to remaining_props before the merge:
# Convert rest-prop keys to camelCase to match component-memo behaviour
camel_remaining = {format.to_camel_case(k): v for k, v in remaining_props.items()}
bound_props.update(camel_remaining)1 participant
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.
All Submissions:
Type of change
Please delete options that are not relevant.
New Feature Submission:
Changes To Core Features:
closes #6186