perf: re-render state readers only when a field they read changes

packages/reflex-base/news/+field-state-subscriptions.performance.md

@@ -0,0 +1 @@
+Components re-render only when a state field they read through a state Var changes, instead of on every change to their state.

packages/reflex-base/src/reflex_base/.templates/web/utils/context-registry.js

@@ -31,18 +31,34 @@ StateStoreContext.displayName = "StateStoreContext";
 /**
  * Subscribe a component to one Python substate in the shared store.
  *
+ * With ``fields``, the component re-renders only when one of those fields
+ * changes, and the returned object may hold stale values for other fields.
+ *
  * @param {string} name - The dotted Python state name.
+ * @param {string[]} [fields] - The substate fields the component reads.
  * @returns {object} The current state value for the named substate.
  */
-export function useStateContext(name) {
+export function useStateContext(name, fields) {
   const store = useContext(StateStoreContext);
-  const [subscribe, getSnapshot] = useMemo(
-    () => [
+  // The compiler passes a new array literal on each render, so key on content.
+  const fieldsKey = fields?.join(",");
+  const [subscribe, getSnapshot] = useMemo(() => {
+    let last;
+    return [
       (listener) => store.subscribe(name, listener),
-      () => store.getSnapshot(name) ?? app.initialState?.[name],
-    ],
-    [store, name],
-  );
+      () => {
+        const next = store.getSnapshot(name) ?? app.initialState?.[name];
+        if (
+          fields !== undefined &&
+          last !== undefined &&
+          fields.every((field) => Object.is(last[field], next[field]))
+        ) {
+          return last;
+        }
+        return (last = next);
+      },
+    ];
+  }, [store, name, fieldsKey]);
   return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
 }
 

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

@@ -18,6 +18,57 @@
     from reflex_base.components.component import Component
 
 
+_STATE_HOOK = re.compile(
+    r'const (\w+) = useStateContext\(("[^"]*")(?:, (\[[^\]]*\]))?\)'
+)
+
+
+def _merge_state_hooks(
+    hooks: dict[str, VarData | None],
+) -> dict[str, VarData | None]:
+    """Merge the field subscriptions to one state into a single hook.
+
+    A hook without a field list reads the whole substate, so it wins over the
+    field lists for the same state.
+
+    Args:
+        hooks: The hooks of one component.
+
+    Returns:
+        The hooks with one ``useStateContext`` call per state.
+    """
+    fields_by_name: dict[str, dict[str, None] | None] = {}
+    state_hooks = 0
+    for hook in hooks:
+        if "useStateContext(" not in hook or not (match := _STATE_HOOK.fullmatch(hook)):
+            continue
+        state_hooks += 1
+        var_name, _, field_list = match.groups()
+        fields = fields_by_name.get(var_name, {})
+        if fields is not None and field_list is not None:
+            fields.update(dict.fromkeys(json.loads(field_list)))
+            fields_by_name[var_name] = fields
+        else:
+            fields_by_name[var_name] = None
+    if state_hooks == len(fields_by_name):
+        return hooks
+    merged = {}
+    for hook, data in hooks.items():
+        match = _STATE_HOOK.fullmatch(hook) if "useStateContext(" in hook else None
+        if match is None:
+            merged[hook] = data
+            continue
+        var_name, state_name, _ = match.groups()
+        if var_name not in fields_by_name:
+            continue
+        fields = fields_by_name.pop(var_name)
+        args = (
+            state_name if fields is None else f"{state_name}, {json.dumps([*fields])}"
+        )
+        merged[f"const {var_name} = useStateContext({args})"] = data
+    return merged
+
+
 def _sort_hooks(
     hooks: dict[str, VarData | None],
 ) -> tuple[list[str], list[str], list[str]]:
@@ -29,6 +80,7 @@ def _sort_hooks(
     Returns:
         The sorted hooks.
     """
+    hooks = _merge_state_hooks(hooks)
     internal_hooks = []
     pre_trigger_hooks = []
     post_trigger_hooks = []

packages/reflex-base/src/reflex_base/vars/base.py

@@ -637,12 +637,13 @@ def from_state(cls, state: type[BaseState] | str, field_name: str = "") -> VarDa
 
         state_name = state if isinstance(state, str) else state.get_full_name()
         state_var_name = format.format_state_name(state_name)
+        hook_args = json.dumps(state_name)
+        if field_name:
+            hook_args += f", {json.dumps([field_name + FIELD_MARKER])}"
         return VarData(
             state=state_name,
             field_name=field_name,
-            hooks={
-                f"const {state_var_name} = useStateContext({json.dumps(state_name)})": None
-            },
+            hooks={f"const {state_var_name} = useStateContext({hook_args})": None},
             imports={
                 f"$/{constants.Dirs.UTILS}/context-registry": [
                     ImportVar(tag="useStateContext")
@@ -1385,8 +1386,15 @@ def _var_set_state(self, state: type[BaseState] | str) -> Self:
         return StateOperation.create(  # pyright: ignore [reportReturnType]
             formatted_state_name,
             self,
+            # The expression can read any key of the substate, so its hook
+            # subscribes to the whole substate rather than a named field.
             _var_data=VarData.merge(
-                VarData.from_state(state, self._js_expr), self._var_data
+                VarData.from_state(state),
+                VarData(
+                    state=state if isinstance(state, str) else state.get_full_name(),
+                    field_name=self._js_expr,
+                ),
+                self._var_data,
             ),
         ).guess_type()
 

tests/units/compiler/test_context_js_template.py

@@ -373,3 +373,65 @@ def test_client_state_store_notifies_only_substate_subscribers(tmp_path: Path):
         "firstNotifications": 0,
         "secondNotifications": 2,
     }
+
+
+@requires_node
+def test_use_state_context_skips_changes_to_unread_fields(tmp_path: Path) -> None:
+    """A field subscription keeps its snapshot until a listed field changes.
+
+    Args:
+        tmp_path: Directory for the generated modules.
+    """
+    registry_path = (
+        Path(__file__).parents[3]
+        / "packages/reflex-base/src/reflex_base/.templates/web/utils/context-registry.js"
+    )
+    (tmp_path / "registry.mjs").write_text(
+        registry_path.read_text().replace('from "react"', 'from "./react.mjs"')
+    )
+    # One component: useMemo keeps its value while the deps are unchanged.
+    (tmp_path / "react.mjs").write_text(
+        """
+export let store;
+export const setStore = (value) => (store = value);
+let memo;
+export const useMemo = (fn, deps) => {
+  if (!memo || deps.some((dep, i) => !Object.is(dep, memo.deps[i]))) {
+    memo = { deps, value: fn() };
+  }
+  return memo.value;
+};
+export const useContext = () => store;
+export const useSyncExternalStore = (subscribe, getSnapshot) => getSnapshot();
+export const createContext = (defaultValue) => ({ defaultValue });
+"""
+    )
+    driver = tmp_path / "driver.mjs"
+    driver.write_text(
+        """
+import * as react from "./react.mjs";
+import { useStateContext } from "./registry.mjs";
+const state = { s: { a: 1, b: 1 } };
+react.setStore({
+  getSnapshot: (name) => state[name],
+  subscribe: () => () => {},
+});
+const first = useStateContext("s", ["a"]);
+state.s = { ...state.s, b: 2 };
+const afterUnread = useStateContext("s", ["a"]);
+state.s = { ...state.s, a: 2 };
+const afterRead = useStateContext("s", ["a"]);
+process.stdout.write(JSON.stringify({
+  unreadKeepsSnapshot: afterUnread === first,
+  readTakesSnapshot: afterRead === state.s,
+}));
+"""
+    )
+    result = subprocess.run(
+        ["node", str(driver)], capture_output=True, encoding="utf-8"
+    )
+    assert result.returncode == 0, result.stderr
+    assert json.loads(result.stdout) == {
+        "unreadKeepsSnapshot": True,
+        "readTakesSnapshot": True,
+    }

tests/units/components/test_component.py

@@ -2171,7 +2171,7 @@ def add_style(self):  # pyright: ignore [reportIncompatibleMethodOverride]
     page._add_style_recursive(Style())
 
     assert (
-        f'const {test_state.get_name()} = useStateContext("{test_state.get_full_name()}")'
+        f'const {test_state.get_name()} = useStateContext("{test_state.get_full_name()}", ["num{FIELD_MARKER}"])'
         in page._get_all_hooks_internal()
     )
     assert any(

tests/units/reflex_base/compiler/test_templates.py

@@ -0,0 +1,32 @@
+"""Tests for reflex_base.compiler.templates."""
+
+from reflex_base.compiler.templates import _sort_hooks
+from reflex_base.constants import Hooks
+from reflex_base.vars.base import VarData
+
+_INTERNAL = VarData(position=Hooks.HookPosition.INTERNAL)
+
+
+def test_sort_hooks_merges_field_subscriptions_per_state():
+    """Each state gets one hook that lists every field the component reads."""
+    internal, _, _ = _sort_hooks({
+        'const s = useStateContext("s", ["a"])': _INTERNAL,
+        'const t = useStateContext("t", ["x"])': _INTERNAL,
+        'const s = useStateContext("s", ["b", "a"])': _INTERNAL,
+        "const other = useOther()": _INTERNAL,
+    })
+    assert internal == [
+        'const s = useStateContext("s", ["a", "b"])',
+        'const t = useStateContext("t", ["x"])',
+        "const other = useOther()",
+    ]
+
+
+def test_sort_hooks_whole_substate_subscription_wins():
+    """A hook without a field list reads the whole substate, in either order."""
+    for hooks in (
+        ['const s = useStateContext("s", ["a"])', 'const s = useStateContext("s")'],
+        ['const s = useStateContext("s")', 'const s = useStateContext("s", ["a"])'],
+    ):
+        internal, _, _ = _sort_hooks(dict.fromkeys(hooks, _INTERNAL))
+        assert internal == ['const s = useStateContext("s")']

tests/units/reflex_base/vars/test_base.py

@@ -22,6 +22,7 @@
 import pytest
 from reflex_base import constants
 from reflex_base.constants import RouteArgType
+from reflex_base.constants.state import FIELD_MARKER
 from reflex_base.environment import _load_dotenv_from_files, environment
 from reflex_base.utils import serializers
 from reflex_base.utils.exceptions import (
@@ -2189,3 +2190,25 @@ class Unannotated(Parent):
         assert child_field.default == 2
         assert isinstance(vars(child)["__counter__"], Field)
         assert child().__counter__ == 2
+
+
+def test_from_state_hook_subscribes_to_the_field():
+    """A state field's hook names the field, so its readers skip other changes."""
+    assert VarData.from_state("s", "a").hooks == (
+        f'const s = useStateContext("s", ["a{FIELD_MARKER}"])',
+    )
+    assert VarData.from_state("s").hooks == ('const s = useStateContext("s")',)
+
+
+@pytest.mark.parametrize("js_expr", ["value", "items_rx_state_[0]"])
+def test_var_set_state_subscribes_to_the_whole_substate(js_expr: str):
+    """The expression may read any key, so its only hook reads the whole substate.
+
+    Args:
+        js_expr: The expression read from the state.
+    """
+    var = Var(_js_expr=js_expr, _var_type=str)._var_set_state("s")
+    var_data = var._get_all_var_data()
+    assert var_data is not None
+    assert var_data.hooks == ('const s = useStateContext("s")',)
+    assert var_data.field_dependencies == {"s": (js_expr,)}