fix: flatten providers for large state trees

.github/workflows/examples.yml

@@ -82,7 +82,7 @@ jobs:
       - name: Install reflex 0.8.23 into a fresh venv
         run: |
           uv venv --python 3.12 "$RUNNER_TEMP/reflex-0.8.23"
-          uv pip install --python "$RUNNER_TEMP/reflex-0.8.23/bin/python" "reflex==0.8.23" "plotly>=6.0,<7.0"
+          uv pip install --python "$RUNNER_TEMP/reflex-0.8.23/bin/python" "reflex==0.8.23" "plotly>=6.0,<7.0" greenlet
           echo "$RUNNER_TEMP/reflex-0.8.23/bin" >> "$GITHUB_PATH"
       - name: Show the reflex under test
         working-directory: examples/playground

docs/app/tests/frontend_quality.test.mjs

@@ -147,7 +147,7 @@ with RegistrationContext(), patch('reflex.compiler.compiler.get_config', return_
     bundle_library('quality-lazy-fixture')
     print(compile_app_root(rx.el.div('Quality'))[1])
 `], {cwd: root, encoding: "utf8"});
-  const setup = new Function("window", "React", "emotion_react", "utils_context", "utils_state", "loadFixture",
+  const setup = new Function("window", "React", "emotion_react", "utils_context", "utils_context_registry", "utils_state", "loadFixture",
     rootCode.replace(/^import .*$/gm, "")
       .replace(/^export default function/gm, "function")
       .replace(/^export /gm, "")
@@ -165,7 +165,7 @@ with RegistrationContext(), patch('reflex.compiler.compiler.get_config', return_
   try {
     let calls = 0;
     globalThis.window = {};
-    setup(window, react, {}, {}, {}, async () => {calls++; return {value: 42};});
+    setup(window, react, {}, {}, {}, {}, async () => {calls++; return {value: 42};});
     assert.equal(calls, 0);
     assert.equal(window.__reflex.react, react);
     assert.equal(window.__reflex["quality-lazy-fixture"], undefined);
@@ -179,7 +179,7 @@ with RegistrationContext(), patch('reflex.compiler.compiler.get_config', return_
 
     let attempts = 0;
     globalThis.window = {};
-    setup(window, react, {}, {}, {}, async () => {
+    setup(window, react, {}, {}, {}, {}, async () => {
       if (++attempts === 1) throw new Error("Temporary download failure");
       return {value: 43};
     });

news/7441.performance.md

@@ -0,0 +1 @@
+Use one shared state-store context with per-substate subscriptions to keep client and server render trees shallow for apps with many state subclasses.

packages/reflex-base/news/+large-state-provider-tree.performance.md

@@ -0,0 +1 @@
+Use one shared state-store context with per-substate subscriptions to keep client and server render trees shallow for apps with many state subclasses.

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

@@ -4,7 +4,12 @@
 // a module the compiler rewrites: Vite re-executes a rewritten module on hot
 // update, and every module it imports along with it, but this module imports
 // nothing that can change and is never re-executed.
-import { createContext } from "react";
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useSyncExternalStore,
+} from "react";
 
 export const ColorModeContext = createContext({
   rawColorMode: "system",
@@ -15,33 +20,61 @@ export const ColorModeContext = createContext({
 export const UploadFilesContext = createContext(null);
 export const DispatchContext = createContext(null);
 export const EventLoopContext = createContext(null);
+export const StateStoreContext = createContext(null);
+const stateContexts = new Map();
 
 ColorModeContext.displayName = "ColorModeContext";
 UploadFilesContext.displayName = "UploadFilesContext";
 DispatchContext.displayName = "DispatchContext";
 EventLoopContext.displayName = "EventLoopContext";
-
-const stateContexts = new Map();
+StateStoreContext.displayName = "StateStoreContext";
 
 /**
- * Return the context object carrying the named Python state.
- *
- * The first call for a name creates the context; later calls, including ones
- * made from a re-executed generated module, return the same object.
+ * Return the stable legacy context for a named Python state.
  *
  * @param {string} name - The dotted Python state name.
- * @returns {React.Context} The state's context object.
+ * @returns {React.Context} The named state's context object.
  */
 export function getStateContext(name) {
   let context = stateContexts.get(name);
   if (context === undefined) {
     context = createContext(null);
     context.displayName = `StateContext(${name})`;
+    context.stateName = name;
     stateContexts.set(name, context);
   }
   return context;
 }
 
+/**
+ * Subscribe a component to one Python substate in the shared store.
+ *
+ * @param {string} name - The dotted Python state name.
+ * @returns {object} The current state value for the named substate.
+ */
+export function useStateContext(name) {
+  const store = useContext(StateStoreContext);
+  const subscribe = useCallback(
+    (listener) => store.subscribe(name, listener),
+    [store, name],
+  );
+  const getSnapshot = useCallback(
+    () => store.getSnapshot(name) ?? app.initialState?.[name],
+    [store, name],
+  );
+  return useSyncExternalStore(subscribe, getSnapshot, getSnapshot);
+}
+
+/**
+ * Read a state through a legacy StateContexts entry using the shared store.
+ *
+ * @param {React.Context} context - The named state's stable context handle.
+ * @returns {object} The current state value.
+ */
+export function useLegacyStateContext(context) {
+  return useStateContext(context.stateName);
+}
+
 // Values the generated context module registers each time it executes. The
 // static runtime reads them through this object at call time instead of
 // importing the generated module, so a compile never re-executes the runtime.

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

@@ -344,7 +344,6 @@ def context_template(
         f'{format_state_name(state_name)}: getStateContext("{state_name}"),'
         for state_name in initial_state
     ])
-
     state_str = (
         rf"""
 export const state_name = "{state_name}"
@@ -411,11 +410,6 @@ def context_template(
 """
     )
 
-    substates_str = "".join(
-        f"\n  ['{state_name}', '{format_state_name(state_name)}'],"
-        for state_name in initial_state
-    )
-
     disable_owner_stacks_str = (
         r"""
 // Disable React dev-build owner-stack capture: the per-element Error()
@@ -444,12 +438,12 @@ def context_template(
         else ""
     )
 
-    return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useReducer, useRef, useState, createElement, useEffect, useLayoutEffect }} from "react"
+    return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useState, createElement, useEffect }} from "react"
 import {{ applyDelta, ReflexEvent, hydrateClientStorage, useEventLoop, refs }} from "$/utils/state"
-import {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, getStateContext, registerApp, eventLoop }} from "$/utils/context-registry"
+import {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, StateStoreContext, getStateContext, useStateContext, registerApp, eventLoop }} from "$/utils/context-registry"
 import {{ jsx }} from "@emotion/react";
 {disable_owner_stacks_str}
-export {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext }};
+export {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, StateStoreContext }};
 export const initialState = {initial_state_json}
 export const initialStateHashes = {"[]" if not initial_state_hashes else json_dumps(initial_state_hashes)}
 
@@ -536,78 +530,89 @@ def context_template(
   );
 }}
 
-// ``useLayoutEffect`` warns when rendered on the server, where no effect runs
-// at all, so fall back to ``useEffect`` there.
-const useIsomorphicLayoutEffect =
-  typeof document !== "undefined" ? useLayoutEffect : useEffect;
-
-// Holds the mutable substate -> dispatch registry that ``SubstateProvider``
-// writes into and ``EventLoopProvider`` reads. The registry object identity is
-// stable for the lifetime of the tree, so neither adding a dispatcher nor
-// updating a substate re-renders the consumers of ``DispatchContext``.
-const DispatchProvider = ({{ children }}) => {{
-  const dispatchers = useRef({{}});
-  return useMemo(
-    () =>
-      createElement(DispatchContext, {{ value: dispatchers.current }}, children),
-    [children],
-  );
-}};
-
-// One provider per substate: each owns its own reducer, so a delta for one
-// substate only re-renders its provider instead of recreating every provider.
-const SubstateProvider = ({{ children, substateName, contextName }}) => {{
-  const dispatchers = useContext(DispatchContext);
-  const [state, dispatchSubstate] = useReducer(
-    applyDelta,
-    initialState[substateName],
-  );
-  // A layout effect, not a passive one: layout effects for the whole commit
-  // run before any passive effect, so every dispatcher is registered before
-  // ``EventLoopProvider`` (mounted below this provider) connects the socket.
-  // A delta naming an unregistered substate is skipped with a warning.
-  useIsomorphicLayoutEffect(() => {{
-    dispatchers[substateName] = dispatchSubstate;
-    return () => {{
-      delete dispatchers[substateName];
-    }};
-  }}, [dispatchers, dispatchSubstate, substateName]);
-  return useMemo(
-    () => createElement(StateContexts[contextName], {{ value: state }}, children),
-    [children, state, contextName],
+// Dispatch updates happen outside render; useSyncExternalStore keeps reads
+// consistent while notifying only consumers of the changed substate.
+function createStateStore() {{
+  const state = {{ ...initialState }};
+  const listeners = new Map();
+  const dispatchers = {{}};
+  const registerSubstate = (substateName, initialValue) => {{
+    if (!listeners.has(substateName)) {{
+      listeners.set(substateName, new Set());
+    }}
+    if (!Object.hasOwn(state, substateName)) {{
+      state[substateName] = initialValue;
+    }}
+    if (!Object.hasOwn(dispatchers, substateName)) {{
+      dispatchers[substateName] = (delta) => {{
+        const currentState = state[substateName];
+        if (
+          Object.keys(delta).every((key) =>
+            Object.is(currentState[key], delta[key]),
+          )
+        ) {{
+          return;
+        }}
+        const nextState = applyDelta(currentState, delta);
+        state[substateName] = nextState;
+        listeners.get(substateName)?.forEach((listener) => listener());
+      }};
+    }}
+  }};
+  Object.entries(state).forEach(([substateName, value]) =>
+    registerSubstate(substateName, value),
   );
-}};
 
-// ``[substateName, contextName]`` for every substate, outermost first.
-const SUBSTATES = [{substates_str}
-];
+  return {{
+    dispatchers,
+    syncInitialState: (nextInitialState) => {{
+      Object.entries(nextInitialState).forEach(([substateName, value]) => {{
+        const isNewSubstate = !Object.hasOwn(state, substateName);
+        registerSubstate(substateName, value);
+        const currentState = state[substateName];
+        const addedFields = Object.fromEntries(
+          Object.entries(value).filter(([key]) => !Object.hasOwn(currentState, key)),
+        );
+        if (Object.keys(addedFields).length > 0) {{
+          state[substateName] = {{ ...currentState, ...addedFields }};
+        }}
+        if (isNewSubstate || Object.keys(addedFields).length > 0) {{
+          listeners.get(substateName)?.forEach((listener) => listener());
+        }}
+      }});
+    }},
+    getSnapshot: (substateName) => state[substateName],
+    subscribe: (substateName, listener) => {{
+      let substateListeners = listeners.get(substateName);
+      if (!substateListeners) {{
+        substateListeners = new Set();
+        listeners.set(substateName, substateListeners);
+      }}
+      substateListeners.add(listener);
+      return () => substateListeners.delete(listener);
+    }},
+  }};
+}}
+
+function StateStoreProviders({{ children, store }}) {{
+  return createElement(
+    StateStoreContext,
+    {{ value: store }},
+    createElement(DispatchContext, {{ value: store.dispatchers }}, children),
+  );
+}}
 
 function ClientStateProvider({{ children }}) {{
-  return useMemo(() => {{
-    let tree = children;
-    for (let i = SUBSTATES.length - 1; i >= 0; i--) {{
-      const [substateName, contextName] = SUBSTATES[i];
-      tree = createElement(SubstateProvider, {{ substateName, contextName }}, tree);
-    }}
-    return createElement(DispatchProvider, {{}}, tree);
-  }}, [children]);
+  const [store] = useState(createStateStore);
+  useEffect(() => store.syncInitialState(initialState), [store, initialState]);
+  return createElement(StateStoreProviders, {{ children, store }});
 }}
 
-// The server renders once and never applies a delta, so it provides the
-// initial state through bare context providers. ``SubstateProvider`` would add
-// a second render level per substate, and the server renderer recurses once per
-// level, so with many substates rendering a page can exhaust the stack.
 function ServerStateProvider({{ children }}) {{
-  let tree = children;
-  for (let i = SUBSTATES.length - 1; i >= 0; i--) {{
-    const [substateName, contextName] = SUBSTATES[i];
-    tree = createElement(
-      StateContexts[contextName],
-      {{ value: initialState[substateName] }},
-      tree,
-    );
-  }}
-  return createElement(DispatchContext, {{ value: {{}} }}, tree);
+  return createElement(
+    StateStoreProviders,
+    {{ children, store: createStateStore() }},
+  );
 }}
 
 export const StateProvider =

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

@@ -32,6 +32,7 @@ def _default_bundled_libraries() -> list[str]:
         "react",
         "@emotion/react",
         f"$/{constants.Dirs.CONTEXTS_PATH}",
+        f"$/{constants.Dirs.UTILS}/context-registry",
         f"$/{constants.Dirs.STATE_PATH}",
     ]
 

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

@@ -636,19 +636,19 @@ def from_state(cls, state: type[BaseState] | str, field_name: str = "") -> VarDa
         from reflex_base.utils import format
 
         state_name = state if isinstance(state, str) else state.get_full_name()
+        state_var_name = format.format_state_name(state_name)
         return VarData(
             state=state_name,
             field_name=field_name,
             hooks={
-                "const {0} = useContext(StateContexts.{0})".format(
-                    format.format_state_name(state_name)
-                ): None
+                f"const {state_var_name} = useStateContext({json.dumps(state_name)})": None
             },
             imports={
-                f"$/{constants.Dirs.CONTEXTS_PATH}": [ImportVar(tag="StateContexts")],
-                "react": [ImportVar(tag="useContext")],
+                f"$/{constants.Dirs.UTILS}/context-registry": [
+                    ImportVar(tag="useStateContext")
+                ],
             },
-            # State Vars read ``StateContexts``/``EventLoopContext``, so the
+            # State Vars read the shared store/``EventLoopContext``, so the
             # providers must enclose every component that uses them.
             app_wraps=get_event_app_wraps(),
         )