From 4e5cd02fb968aaf3851c890ea075cac811f9dd3e Mon Sep 17 00:00:00 2001 From: Farhan Date: Fri, 16 Oct 2026 22:28:38 +0500 Subject: [PATCH 1/2] perf: re-render state readers only when a field they read changes Each state var hook now names its field, the compiler merges the hooks of a component into one useStateContext call per state with the union of fields, and useStateContext keeps the previous snapshot while those fields are unchanged. A hook without a field list still reads the whole substate. --- .../+field-state-subscriptions.performance.md | 1 + .../.templates/web/utils/context-registry.js | 30 ++++++--- .../src/reflex_base/compiler/templates.py | 52 ++++++++++++++++ .../reflex-base/src/reflex_base/vars/base.py | 13 ++-- .../compiler/test_context_js_template.py | 62 +++++++++++++++++++ tests/units/components/test_component.py | 2 +- tests/units/reflex_base/compiler/__init__.py | 0 .../reflex_base/compiler/test_templates.py | 32 ++++++++++ tests/units/reflex_base/vars/test_base.py | 17 +++++ 9 files changed, 197 insertions(+), 12 deletions(-) create mode 100644 packages/reflex-base/news/+field-state-subscriptions.performance.md create mode 100644 tests/units/reflex_base/compiler/__init__.py create mode 100644 tests/units/reflex_base/compiler/test_templates.py diff --git a/packages/reflex-base/news/+field-state-subscriptions.performance.md b/packages/reflex-base/news/+field-state-subscriptions.performance.md new file mode 100644 index 00000000000..e6ebe226a81 --- /dev/null +++ b/packages/reflex-base/news/+field-state-subscriptions.performance.md @@ -0,0 +1 @@ +Components re-render only when a state field they read changes, instead of on every change to their state. diff --git a/packages/reflex-base/src/reflex_base/.templates/web/utils/context-registry.js b/packages/reflex-base/src/reflex_base/.templates/web/utils/context-registry.js index 06778dc0092..d4909b46318 100644 --- a/packages/reflex-base/src/reflex_base/.templates/web/utils/context-registry.js +++ b/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); } diff --git a/packages/reflex-base/src/reflex_base/compiler/templates.py b/packages/reflex-base/src/reflex_base/compiler/templates.py index 0c65b8b73c8..8eea6cb4145 100644 --- a/packages/reflex-base/src/reflex_base/compiler/templates.py +++ b/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 = [] diff --git a/packages/reflex-base/src/reflex_base/vars/base.py b/packages/reflex-base/src/reflex_base/vars/base.py index 0671d80c917..4e7760faac3 100644 --- a/packages/reflex-base/src/reflex_base/vars/base.py +++ b/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,12 @@ 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 the plain + # from_state hook subscribes the component to the whole substate. _var_data=VarData.merge( - VarData.from_state(state, self._js_expr), self._var_data + VarData.from_state(state, self._js_expr), + VarData.from_state(state), + self._var_data, ), ).guess_type() diff --git a/tests/units/compiler/test_context_js_template.py b/tests/units/compiler/test_context_js_template.py index f0c6a9511c5..86e455c6b33 100644 --- a/tests/units/compiler/test_context_js_template.py +++ b/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, + } diff --git a/tests/units/components/test_component.py b/tests/units/components/test_component.py index dcba084aa39..684e586c43d 100644 --- a/tests/units/components/test_component.py +++ b/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( diff --git a/tests/units/reflex_base/compiler/__init__.py b/tests/units/reflex_base/compiler/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/units/reflex_base/compiler/test_templates.py b/tests/units/reflex_base/compiler/test_templates.py new file mode 100644 index 00000000000..8f9ddb94740 --- /dev/null +++ b/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")'] diff --git a/tests/units/reflex_base/vars/test_base.py b/tests/units/reflex_base/vars/test_base.py index 58517ddaa52..12aefe2e684 100644 --- a/tests/units/reflex_base/vars/test_base.py +++ b/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,19 @@ 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")',) + + +def test_var_set_state_subscribes_to_the_whole_substate(): + """The expression may read any key, so the whole-substate hook must be present.""" + var = Var(_js_expr="value", _var_type=str)._var_set_state("s") + var_data = var._get_all_var_data() + assert var_data is not None + assert 'const s = useStateContext("s")' in var_data.hooks From dae7365a8cfd8165ddfa63b5d633a99d7b613d30 Mon Sep 17 00:00:00 2001 From: Farhan Date: Fri, 16 Oct 2026 22:35:10 +0500 Subject: [PATCH 2/2] fix: give _var_set_state only a whole-substate hook Its expression is not a field name, so a field hook built from it could hold a ']' the hook merge cannot parse, leaving two declarations of the state. --- .../news/+field-state-subscriptions.performance.md | 2 +- packages/reflex-base/src/reflex_base/vars/base.py | 9 ++++++--- tests/units/reflex_base/vars/test_base.py | 14 ++++++++++---- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/packages/reflex-base/news/+field-state-subscriptions.performance.md b/packages/reflex-base/news/+field-state-subscriptions.performance.md index e6ebe226a81..fd9fcecc78d 100644 --- a/packages/reflex-base/news/+field-state-subscriptions.performance.md +++ b/packages/reflex-base/news/+field-state-subscriptions.performance.md @@ -1 +1 @@ -Components re-render only when a state field they read changes, instead of on every change to their state. +Components re-render only when a state field they read through a state Var changes, instead of on every change to their state. diff --git a/packages/reflex-base/src/reflex_base/vars/base.py b/packages/reflex-base/src/reflex_base/vars/base.py index 4e7760faac3..4804de56786 100644 --- a/packages/reflex-base/src/reflex_base/vars/base.py +++ b/packages/reflex-base/src/reflex_base/vars/base.py @@ -1386,11 +1386,14 @@ 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 the plain - # from_state hook subscribes the component to the whole substate. + # 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), 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() diff --git a/tests/units/reflex_base/vars/test_base.py b/tests/units/reflex_base/vars/test_base.py index 12aefe2e684..37856bb753b 100644 --- a/tests/units/reflex_base/vars/test_base.py +++ b/tests/units/reflex_base/vars/test_base.py @@ -2200,9 +2200,15 @@ def test_from_state_hook_subscribes_to_the_field(): assert VarData.from_state("s").hooks == ('const s = useStateContext("s")',) -def test_var_set_state_subscribes_to_the_whole_substate(): - """The expression may read any key, so the whole-substate hook must be present.""" - var = Var(_js_expr="value", _var_type=str)._var_set_state("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 'const s = useStateContext("s")' in var_data.hooks + assert var_data.hooks == ('const s = useStateContext("s")',) + assert var_data.field_dependencies == {"s": (js_expr,)}