Reduce MutableProxy per-element overhead and cache per-field proxies

news/6740.performance.md

@@ -0,0 +1 @@
+Speed up MutableProxy: immutable elements retrieved through a proxy skip the dataclasses frame-walk check, and the proxy for each mutable state var is cached per instance instead of rebuilt on every attribute read.

packages/reflex-base/news/6740.performance.md

@@ -0,0 +1 @@
+Reserve the internal `_mutable_proxy_cache` state field name so user vars cannot collide with the per-instance proxy cache.

packages/reflex-base/src/reflex_base/utils/types.py

@@ -161,7 +161,13 @@ def __call__(
     dict: Dict,  # noqa: UP006
 }
 
-RESERVED_BACKEND_VAR_NAMES = {"_abc_impl", "_backend_vars", "_was_touched", "_mixin"}
+RESERVED_BACKEND_VAR_NAMES = {
+    "_abc_impl",
+    "_backend_vars",
+    "_was_touched",
+    "_mixin",
+    "_mutable_proxy_cache",
+}
 
 
 class Unset:

reflex/istate/proxy.py

@@ -522,22 +522,24 @@ def _wrap_recursive(self, value: Any) -> Any:
         Returns:
             The wrapped value.
         """
-        # When called from dataclasses internal code, return the unwrapped value
-        if self._is_called_from_dataclasses_internal():
-            return value
         # If we already have a proxy, unwrap and rewrap to make sure the state
         # reference is up to date.
         if isinstance(value, MutableProxy):
             value = value.__wrapped__
+        # Immutable values (the common case when iterating a container of
+        # scalars) never need wrapping nor the frame inspection below.
+        if not is_mutable_type(type(value)):
+            return value
+        # When called from dataclasses internal code, return the unwrapped value
+        if self._is_called_from_dataclasses_internal():
+            return value
         # Recursively wrap mutable types.
-        if is_mutable_type(type(value)):
-            base_cls = globals()[self.__base_proxy__]
-            return base_cls(
-                wrapped=value,
-                state=self._self_state,
-                field_name=self._self_field_name,
-            )
-        return value
+        base_cls = globals()[self.__base_proxy__]
+        return base_cls(
+            wrapped=value,
+            state=self._self_state,
+            field_name=self._self_field_name,
+        )
 
     def _wrap_recursive_decorator(
         self, wrapped: Callable, instance: BaseState, args: list, kwargs: dict

reflex/state.py

@@ -438,6 +438,12 @@ class BaseState(EvenMoreBasicBaseState):
     # Whether the state has ever been touched since instantiation.
     _was_touched: bool = field(default=False, is_var=False)
 
+    # Per-field cache of the MutableProxy wrapping each mutable var, so
+    # repeated reads don't rebuild the proxy. Never pickled.
+    _mutable_proxy_cache: builtins.dict[str, MutableProxy] = field(
+        default_factory=builtins.dict, is_var=False
+    )
+
     # A special event handler for setting base vars.
     setvar: ClassVar[EventHandler]
 
@@ -1478,7 +1484,12 @@ def _get_attribute(self, name: str) -> Any:
             name in super().__getattribute__("base_vars") or name in backend_vars
         ):
             # track changes in mutable containers (list, dict, set, etc)
-            return MutableProxy(wrapped=value, state=self, field_name=name)
+            cache = super().__getattribute__("_mutable_proxy_cache")
+            proxy = cache.get(name)
+            if proxy is None or proxy.__wrapped__ is not value:
+                proxy = MutableProxy(wrapped=value, state=self, field_name=name)
+                cache[name] = proxy
+            return proxy
 
         return value
 
@@ -1513,6 +1524,8 @@ def __setattr__(self, name: str, value: Any):
 
         if name in self.backend_vars:
             self._backend_vars.__setitem__(name, value)
+            # Drop the proxy wrapping the replaced value.
+            self.__dict__["_mutable_proxy_cache"].pop(name, None)
             self.dirty_vars.add(name)
             self._mark_dirty()
             return
@@ -1544,6 +1557,9 @@ def __setattr__(self, name: str, value: Any):
         # Set the attribute.
         object.__setattr__(self, name, value)
 
+        # Drop the proxy wrapping the replaced value.
+        self.__dict__["_mutable_proxy_cache"].pop(name, None)
+
         # Add the var to the dirty list.
         if name in self.base_vars:
             self.dirty_vars.add(name)
@@ -2069,6 +2085,8 @@ def __getstate__(self):
         state.pop("parent_state", None)
         state.pop("substates", None)
         state.pop("_was_touched", None)
+        # Proxies wrap live state references and are rebuilt on access.
+        state.pop("_mutable_proxy_cache", None)
         # Remove all inherited vars.
         for inherited_var_name in self.inherited_vars:
             state.pop(inherited_var_name, None)
@@ -2084,6 +2102,8 @@ def __setstate__(self, state: dict[str, Any]):
         """
         state["parent_state"] = None
         state["substates"] = {}
+        # The proxy cache is never pickled; recreate it on the restored instance.
+        state.setdefault("_mutable_proxy_cache", {})
         for key, value in state.items():
             object.__setattr__(self, key, value)
 

tests/units/istate/test_proxy.py

@@ -71,3 +71,51 @@ async def mock_modify_state_context(*args, **kwargs):
     # After the exception, we should be able to enter the context again without issues
     async with state_proxy:
         pass
+
+
+def test_mutable_proxy_cached_per_field():
+    """Repeated reads of a mutable var reuse the proxy until reassignment."""
+    state = ProxyTestState()
+    first = state.items
+    assert isinstance(first, MutableProxy)
+    assert state.items is first
+    # Reassignment invalidates the cached proxy.
+    state.items = [Item(2)]
+    second = state.items
+    assert isinstance(second, MutableProxy)
+    assert second is not first
+    assert second[0].id == 2
+    # In-place mutation keeps the same wrapped object, so the proxy is reused.
+    second.append(Item(3))
+    assert state.items is second
+    # Reassignment immediately evicts the cache entry, so no strong reference
+    # to the replaced value lingers until the next read.
+    state.items = [Item(4)]
+    assert "items" not in state.__dict__["_mutable_proxy_cache"]
+
+
+def test_mutable_proxy_cache_not_serialized():
+    """The per-instance proxy cache never leaks into pickles or copies."""
+    state = ProxyTestState()
+    state.items.append(Item(1))  # populate the proxy cache
+    assert state.__dict__["_mutable_proxy_cache"]
+    assert "_mutable_proxy_cache" not in state.__getstate__()
+
+    restored = pickle.loads(pickle.dumps(state))
+    # The cache is recreated empty on the restored instance.
+    assert restored.__dict__["_mutable_proxy_cache"] == {}
+    restored_items = restored.items
+    assert isinstance(restored_items, MutableProxy)
+    # The restored proxy tracks the restored state, not the original.
+    assert restored_items._self_state is restored
+
+
+def test_mutable_proxy_iteration_yields_plain_immutables():
+    """Iterating a proxied container returns immutable elements unwrapped."""
+    state = ProxyTestState()
+    state.items = [Item(1), Item(2)]
+    numbers = [item.id for item in state.items]
+    assert numbers == [1, 2]
+    assert all(type(n) is int for n in numbers)
+    # Mutable elements remain wrapped so nested mutations mark the state dirty.
+    assert all(isinstance(item, MutableProxy) for item in state.items)

tests/units/test_state.py

@@ -3243,11 +3243,11 @@ class BaseFieldSetterState(BaseState):
     bfss.dirty_vars.clear()
     assert "c1" not in bfss.dirty_vars
 
-    # Assert identity of MutableProxy
+    # Repeated reads reuse the cached MutableProxy for the same field.
     mp = bfss.c1
     assert isinstance(mp, MutableProxy)
     mp3 = bfss.c1
-    assert mp is not mp3
+    assert mp is mp3
     # Since none of these set calls had values, the state should not be dirty
     assert not bfss.dirty_vars