refactor: freeze components after create and use copy_with for mutations

packages/reflex-base/src/reflex_base/components/component.py

@@ -18,7 +18,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, TypeVar, cast, get_args, get_origin
 
 from rich.markup import escape
-from typing_extensions import dataclass_transform
+from typing_extensions import Self, dataclass_transform
 
 from reflex_base import constants
 from reflex_base.breakpoints import Breakpoints
@@ -266,9 +266,20 @@ class BaseComponent(metaclass=BaseComponentMeta):
     This is something that can be rendered as a Component via the Reflex compiler.
     """
 
-    children: list[BaseComponent] = field(
+    _frozen: ClassVar[bool] = False
+
+    # Render-path caches; allowed to be written even on frozen instances.
+    _CACHE_ATTRS: ClassVar[frozenset[str]] = frozenset({
+        "_cached_render_result",
+        "_vars_cache",
+        "_imports_cache",
+        "_hooks_internal_cache",
+        "_get_component_prop_property",
+    })
+
+    children: tuple[BaseComponent, ...] = field(
         doc="The children nested within the component.",
-        default_factory=list,
+        default_factory=tuple,
         is_javascript_property=False,
     )
 
@@ -293,24 +304,73 @@ def __init__(
         Args:
             **kwargs: The kwargs to pass to the component.
         """
+        if "children" in kwargs:
+            kwargs["children"] = tuple(kwargs["children"])
         for key, value in kwargs.items():
             setattr(self, key, value)
         for name, value in self.get_fields().items():
             if name not in kwargs:
                 setattr(self, name, value.default_value())
 
+    def __setattr__(self, key: str, value: Any) -> None:
+        """Block writes to frozen components, except for cache attributes.
+
+        Args:
+            key: The attribute name.
+            value: The attribute value.
+
+        Raises:
+            AttributeError: If the component is frozen and the attribute is not a cache.
+        """
+        if self.__dict__.get("_frozen", False) and key not in type(self)._CACHE_ATTRS:
+            msg = (
+                f"Cannot set {key!r} on frozen {type(self).__name__}; "
+                "use copy_with() to create a modified copy."
+            )
+            raise AttributeError(msg)
+        super().__setattr__(key, value)
+
+    def _freeze(self) -> None:
+        """Mark this component as frozen.
+
+        Subsequent attribute writes outside the cache allowlist will raise.
+        """
+        object.__setattr__(self, "_frozen", True)
+
+    def copy_with(self, **updates: Any) -> Self:
+        """Return a frozen shallow copy with updated fields.
+
+        Bypasses ``__setattr__`` for speed and to skip the freeze guard.
+        Render-path caches are dropped because they may depend on the fields
+        being replaced.
+
+        Args:
+            **updates: Field values to override on the copy.
+
+        Returns:
+            A new frozen instance with the requested updates applied.
+        """
+        new = self.__class__.__new__(self.__class__)
+        d = vars(new)
+        d.update(vars(self))
+        for cache_attr in type(self)._CACHE_ATTRS:
+            d.pop(cache_attr, None)
+        if "children" in updates:
+            updates["children"] = tuple(updates["children"])
+        d.update(updates)
+        d["_frozen"] = True
+        return new
+
     def set(self, **kwargs):
-        """Set the component props.
+        """Set the component props, returning a new frozen instance.
 
         Args:
             **kwargs: The kwargs to set.
 
         Returns:
-            The component with the updated props.
+            A new component with the updated props.
         """
-        for key, value in kwargs.items():
-            setattr(self, key, value)
-        return self
+        return self.copy_with(**kwargs)
 
     def __copy__(self) -> BaseComponent:
         """Return a shallow copy suitable for compile-time mutation.
@@ -327,13 +387,7 @@ def __copy__(self) -> BaseComponent:
         new = self.__class__.__new__(self.__class__)
         new_dict = vars(new)
         new_dict.update(vars(self))
-        for attr in (
-            "_cached_render_result",
-            "_vars_cache",
-            "_imports_cache",
-            "_hooks_internal_cache",
-            "_get_component_prop_property",
-        ):
+        for attr in type(self)._CACHE_ATTRS:
             new_dict.pop(attr, None)
         return new
 
@@ -1223,9 +1277,11 @@ def _create(cls: type[T], children: Sequence[BaseComponent], **props: Any) -> T:
         Returns:
             The component.
         """
+        children_tuple = tuple(children)
         comp = cls.__new__(cls)
-        super(Component, comp).__init__(id=props.get("id"), children=list(children))
-        comp._post_init(children=list(children), **props)
+        super(Component, comp).__init__(id=props.get("id"), children=children_tuple)
+        comp._post_init(children=children_tuple, **props)
+        comp._freeze()
         return comp
 
     @classmethod
@@ -1241,10 +1297,12 @@ def _unsafe_create(
         Returns:
             The component.
         """
+        children_tuple = tuple(children)
         comp = cls.__new__(cls)
-        super(Component, comp).__init__(id=props.get("id"), children=list(children))
+        super(Component, comp).__init__(id=props.get("id"), children=children_tuple)
         for prop, value in props.items():
             setattr(comp, prop, value)
+        comp._freeze()
         return comp
 
     def add_style(self) -> dict[str, Any] | None:
@@ -1311,40 +1369,47 @@ def _add_style_recursive(
             theme: The theme to apply. (for retro-compatibility with deprecated _apply_theme API)
 
         Returns:
-            The component with the additional style.
+            A component with the additional style; ``self`` if nothing changed.
 
         Raises:
             UserWarning: If `_add_style` has been overridden.
         """
-        # 1. Default style from `_add_style`/`add_style`.
         if type(self)._add_style != Component._add_style:
             msg = "Do not override _add_style directly. Use add_style instead."
             raise UserWarning(msg)
-        new_style = self._add_style()
-        style_vars = [new_style._var_data]
 
-        # 2. User-defined style from `App.style`.
+        style_addition = self._add_style()
         component_style = self._get_component_style(style)
-        if component_style:
-            new_style.update(component_style)
-            style_vars.append(component_style._var_data)
-
-        # 4. style dict and css props passed to the component instance.
-        new_style.update(self.style)
-        style_vars.append(self.style._var_data)
-
-        new_style._var_data = VarData.merge(*style_vars)
-
-        # Assign the new style
-        self.style = new_style
+        has_style_change = bool(style_addition) or bool(component_style)
 
-        # Recursively add style to the children.
-        for child in self.children:
-            # Skip non-Component children.
+        new_children: list | None = None
+        for i, child in enumerate(self.children):
             if not isinstance(child, Component):
                 continue
-            child._add_style_recursive(style, theme)
-        return self
+            updated = child._add_style_recursive(style, theme)
+            if updated is child:
+                continue
+            if new_children is None:
+                new_children = list(self.children)
+            new_children[i] = updated
+
+        if not has_style_change and new_children is None:
+            return self
+
+        updates: dict[str, Any] = {}
+        if has_style_change:
+            new_style = style_addition
+            style_vars = [new_style._var_data]
+            if component_style:
+                new_style.update(component_style)
+                style_vars.append(component_style._var_data)
+            new_style.update(self.style)
+            style_vars.append(self.style._var_data)
+            new_style._var_data = VarData.merge(*style_vars)
+            updates["style"] = new_style
+        if new_children is not None:
+            updates["children"] = tuple(new_children)
+        return self.copy_with(**updates)
 
     def _get_style(self) -> dict:
         """Get the style for the component.
@@ -2342,8 +2407,7 @@ def get_component(self) -> Component:
         except Exception:
             style = {}
 
-        component._add_style_recursive(style)
-        return component
+        return component._add_style_recursive(style)
 
     def _get_all_app_wrap_components(
         self, *, ignore_ids: set[int] | None = None

packages/reflex-base/src/reflex_base/components/memoize_helpers.py

@@ -154,27 +154,27 @@ def fix_event_triggers_for_memo(
     """Return a component whose event triggers reference memoized ``useCallback``s.
 
     Replaces each (non-lifecycle) event-trigger value with a ``Var`` naming a
-    memoized ``useCallback`` wrapper. The original is never mutated — a
-    page-local clone is taken via ``page_context.own`` on first write.
+    memoized ``useCallback`` wrapper. The original is never mutated — a frozen
+    copy with the rewritten triggers is returned via ``copy_with``.
 
     Args:
         component: The component whose event triggers to memoize.
-        page_context: The active page context, used to obtain a page-local
-            clone before rewriting ``event_triggers``.
+        page_context: The active page context (unused; retained for API
+            compatibility with downstream callers).
 
     Returns:
-        Either ``component`` (when nothing needed rewriting) or a page-local
-        clone with the rewritten ``event_triggers``.
+        Either ``component`` (when nothing needed rewriting) or a new frozen
+        copy with the rewritten ``event_triggers``.
     """
     memo_event_triggers = tuple(get_memoized_event_triggers(component).items())
     if not memo_event_triggers:
         return component
-    owned = page_context.own(component)
-    owned.event_triggers = {
-        **component.event_triggers,
-        **dict(memo_event_triggers),
-    }
-    return owned
+    return component.copy_with(
+        event_triggers={
+            **component.event_triggers,
+            **dict(memo_event_triggers),
+        }
+    )
 
 
 def is_snapshot_boundary(component: Component) -> bool:

packages/reflex-base/src/reflex_base/components/tags/iter_tag.py

@@ -112,6 +112,6 @@ def render_component(self) -> Component:
 
         # Set the component key.
         if component.key is None:
-            component.key = index
+            component = component.copy_with(key=index)
 
         return component

packages/reflex-base/src/reflex_base/plugins/compiler.py

@@ -2,13 +2,12 @@
 
 from __future__ import annotations
 
-import copy
 import dataclasses
 import inspect
 from collections.abc import Callable, Sequence
 from contextvars import ContextVar, Token
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias, TypeVar, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Protocol, TypeAlias, cast
 
 from typing_extensions import Self
 
@@ -32,9 +31,6 @@
     )
 
 
-_BaseComponentT = TypeVar("_BaseComponentT", bound=BaseComponent)
-
-
 class PageDefinition(Protocol):
     """Protocol for page-like objects compiled by :class:`CompileContext`."""
 
@@ -374,8 +370,7 @@ def visit(
                     updated_children = list(children[:index])
                 updated_children.append(compiled_child)
             if updated_children is not None:
-                current_comp = page_context.own(current_comp)
-                current_comp.children = updated_children
+                current_comp = current_comp.copy_with(children=tuple(updated_children))
 
             if isinstance(current_comp, Component):
                 for prop_component in current_comp._get_components_in_props():
@@ -437,8 +432,7 @@ def visit(
                     updated_children = list(children[:index])
                 updated_children.append(compiled_child)
             if updated_children is not None:
-                current_comp = page_context.own(current_comp)
-                current_comp.children = updated_children
+                current_comp = current_comp.copy_with(children=tuple(updated_children))
 
             if isinstance(current_comp, Component):
                 for prop_component in current_comp._get_components_in_props():
@@ -549,8 +543,9 @@ def visit(
             if len(compiled_children) != len(current) or any(
                 a is not b for a, b in zip(compiled_children, current, strict=True)
             ):
-                compiled_component = page_context.own(compiled_component)
-                compiled_component.children = list(compiled_children)
+                compiled_component = compiled_component.copy_with(
+                    children=tuple(compiled_children)
+                )
             return compiled_component
 
         return visit(
@@ -695,38 +690,6 @@ class PageContext(BaseContext):
     # the matching ``leave_component``. Non-empty iff we are inside such a
     # subtree.
     memoize_suppressor_stack: list[int] = dataclasses.field(default_factory=list)
-    # Maps both the user-owned original's ``id()`` and the clone's ``id()`` to
-    # the page-local clone. Lets the walker and plugins rebind children, style,
-    # or event_triggers on a page-local copy without mutating a user-owned
-    # instance that may be referenced from another route.
-    _owned: dict[int, BaseComponent] = dataclasses.field(default_factory=dict)
-    # Strong references to originals keyed by ``id()`` above. Without these,
-    # an original that is only reachable through ``_owned``'s int key can be
-    # garbage collected, and Python may recycle its ``id()`` for a fresh
-    # component, causing ``own()`` to hand back the wrong clone.
-    _owned_refs: list[BaseComponent] = dataclasses.field(default_factory=list)
-
-    def own(self, comp: _BaseComponentT) -> _BaseComponentT:
-        """Return a page-local copy of ``comp``, cloning on first encounter.
-
-        Repeated calls with the same original return the same clone, so
-        mutations from several plugins accumulate on one instance.
-
-        Args:
-            comp: The component the caller is about to mutate.
-
-        Returns:
-            A component the caller may freely mutate without touching any
-            user-owned instance.
-        """
-        existing = self._owned.get(id(comp))
-        if existing is not None:
-            return cast("_BaseComponentT", existing)
-        new = copy.copy(comp)
-        self._owned[id(comp)] = new
-        self._owned[id(new)] = new
-        self._owned_refs.append(comp)
-        return new
 
     def merged_imports(self, *, collapse: bool = False) -> ParsedImportDict:
         """Return the imports accumulated for this page.