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.

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

@@ -1478,7 +1478,17 @@ 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__("__dict__").get("_mutable_proxy_cache")
+            if cache is None:
+                cache = {}
+                object.__setattr__(self, "_mutable_proxy_cache", cache)
+            proxy = cache.get(name)
+            # isinstance also rejects entries degraded by deepcopy, which
+            # copies a MutableProxy as its unwrapped value.
+            if not isinstance(proxy, MutableProxy) or proxy.__wrapped__ is not value:
+                proxy = MutableProxy(wrapped=value, state=self, field_name=name)
+                cache[name] = proxy
+            return proxy
 
         return value
 
@@ -1513,6 +1523,9 @@ def __setattr__(self, name: str, value: Any):
 
         if name in self.backend_vars:
             self._backend_vars.__setitem__(name, value)
+            if (cache := self.__dict__.get("_mutable_proxy_cache")) is not None:
+                # Drop the proxy wrapping the replaced value.
+                cache.pop(name, None)
             self.dirty_vars.add(name)
             self._mark_dirty()
             return
@@ -1544,6 +1557,10 @@ def __setattr__(self, name: str, value: Any):
         # Set the attribute.
         object.__setattr__(self, name, value)
 
+        if (cache := self.__dict__.get("_mutable_proxy_cache")) is not None:
+            # Drop the proxy wrapping the replaced value.
+            cache.pop(name, None)
+
         # Add the var to the dirty list.
         if name in self.base_vars:
             self.dirty_vars.add(name)
@@ -2069,6 +2086,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)

tests/units/istate/test_proxy.py

@@ -71,3 +71,50 @@ 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 "_mutable_proxy_cache" in state.__dict__
+    assert "_mutable_proxy_cache" not in state.__getstate__()
+
+    restored = pickle.loads(pickle.dumps(state))
+    assert "_mutable_proxy_cache" not in restored.__dict__
+    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