Create separate context providing components for each state

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

@@ -0,0 +1 @@
+Reduce browser rendering overhead for state updates, especially in apps with many substates.

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

@@ -388,18 +388,8 @@ def context_template(
 """
     )
 
-    state_reducer_str = "\n".join(
-        rf'const [{format_state_name(state_name)}, dispatch_{format_state_name(state_name)}] = useReducer(applyDelta, initialState["{state_name}"])'
-        for state_name in initial_state
-    )
-
     create_state_contexts_str = "\n".join(
-        rf"createElement(StateContexts.{format_state_name(state_name)},{{value: {format_state_name(state_name)}}},"
-        for state_name in initial_state
-    )
-
-    dispatchers_str = "\n".join(
-        f'"{state_name}": dispatch_{format_state_name(state_name)},'
+        rf"createElement(SubstateProvider, {{substateName: '{state_name}', contextName: '{format_state_name(state_name)}'}},"
         for state_name in initial_state
     )
 
@@ -431,7 +421,7 @@ def context_template(
         else ""
     )
 
-    return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useReducer, useState, createElement, useEffect }} from "react"
+    return rf"""import {"React, " if disable_react_owner_stacks else ""}{{ useContext, useMemo, useReducer, useRef, useState, createElement, useEffect, useLayoutEffect }} from "react"
 import {{ applyDelta, ReflexEvent, hydrateClientStorage, useEventLoop, refs }} from "$/utils/state"
 import {{ ColorModeContext, UploadFilesContext, DispatchContext, EventLoopContext, getStateContext, registerApp, eventLoop }} from "$/utils/context-registry"
 import {{ jsx }} from "@emotion/react";
@@ -518,19 +508,57 @@ def context_template(
   );
 }}
 
-export function StateProvider({{ children }}) {{
-  {state_reducer_str}
-  const dispatchers = useMemo(() => {{
-    return {{
-      {dispatchers_str}
-    }}
-  }}, [])
+// ``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 a fatal state mismatch.
+  useIsomorphicLayoutEffect(() => {{
+    dispatchers[substateName] = dispatchSubstate;
+    return () => {{
+      delete dispatchers[substateName];
+    }};
+  }}, [dispatchers, dispatchSubstate, substateName]);
+  return useMemo(
+    () => createElement(StateContexts[contextName], {{ value: state }}, children),
+    [children, state, contextName],
+  );
+}};
 
-  return (
-    {create_state_contexts_str}
-    createElement(DispatchContext, {{value: dispatchers}}, children)
+export function StateProvider({{ children }}) {{
+  return useMemo(
+    () => (
+    createElement(DispatchProvider, {{}},
+    {create_state_contexts_str}children
     {")" * len(initial_state)}
-  )
+  )),
+    [children],
+  );
 }}"""
 
 

tests/integration/tests_playwright/test_router_query.py

@@ -24,6 +24,8 @@
 
 import pytest
 from playwright.sync_api import Page, expect
+from reflex_base.constants.state import FIELD_MARKER
+from socketio.packet import Packet
 
 from reflex.testing import AppHarness
 
@@ -32,6 +34,11 @@ def RouterQueryApp():
     """App exercising replaceState, redirect, and redirect(replace=True)."""
     import reflex as rx
 
+    class ConnectedState(rx.State):
+        """A sibling updated by the first on-load event after connecting."""
+
+        connected_count: int = 0
+
     class RouterQueryState(rx.State):
         # Incremented by the page on_load handler; proves whether a navigation
         # (and thus on_load) actually fired.
@@ -59,9 +66,11 @@ def name_param(self) -> str:
             return self.router.url.query_parameters.get("name", "")
 
         @rx.event
-        def on_load(self):
-            """Record that the page on_load handler fired."""
+        async def on_load(self):
+            """Update both substates in the same on-load delta."""
             self.load_count += 1
+            connected = await self.get_state(ConnectedState)
+            connected.connected_count += 1
 
         @rx.event
         def replace_via_script(self, query: str):
@@ -141,6 +150,11 @@ def index():
                 read_only=True,
                 id="load-count",
             ),
+            rx.input(
+                value=f"{ConnectedState.connected_count}",
+                read_only=True,
+                id="connected-count",
+            ),
             rx.input(
                 value=f"{RouterQueryState.ping_count}",
                 read_only=True,
@@ -196,6 +210,44 @@ def _load(harness: AppHarness, page: Page) -> str:
     return base
 
 
+def test_initial_connection_dispatches_both_substates(
+    router_query_app: AppHarness, page: Page
+):
+    """The first on-load delta updates both sibling substates in the browser.
+
+    Args:
+        router_query_app: Running application in dev or production mode.
+        page: Browser page receiving the initial connection's state updates.
+    """
+    frames: list[str | bytes] = []
+    page.on(
+        "websocket",
+        lambda socket: socket.on("framereceived", lambda frame: frames.append(frame)),
+    )
+
+    _load(router_query_app, page)
+    expect(page.locator("#connected-count")).to_have_value("1")
+
+    # Socket.IO event packets use the Engine.IO message and event prefixes, 42.
+    for frame in frames:
+        if not isinstance(frame, str) or not frame.startswith("42"):
+            continue
+        packet = Packet(encoded_packet=frame[1:])
+        assert packet.data is not None
+        event, update = packet.data
+        if event != "event":
+            continue
+        updated_substates = sum(
+            fields.get("load_count" + FIELD_MARKER) == 1
+            or fields.get("connected_count" + FIELD_MARKER) == 1
+            for fields in update.get("delta", {}).values()
+        )
+        if updated_substates == 2:
+            break
+    else:
+        pytest.fail("The initial on-load update did not contain both substates")
+
+
 def test_replace_state_is_not_reactive_but_next_event_syncs(
     router_query_app: AppHarness, page: Page
 ):

tests/units/compiler/test_compiler.py

@@ -1669,6 +1669,39 @@ def test_context_template_owner_stack_pin(disable_owner_stacks: bool):
     assert "captureOwnerStack" in rendered
 
 
+def test_context_template_one_provider_per_substate():
+    """Each substate gets its own provider so one delta re-renders one context.
+
+    A single provider owning every reducer means any delta recreates every
+    ``StateContexts`` element; nesting one ``SubstateProvider`` per substate
+    keeps the untouched providers memoized.
+    """
+    from reflex_base.compiler.templates import context_template
+
+    rendered = context_template(
+        is_dev_mode=True,
+        default_color_mode='"light"',
+        initial_state={
+            "reflex___state____state": {},
+            "reflex___state____state__sub": {},
+        },
+        state_name="reflex___state____state",
+    )
+
+    assert (
+        "createElement(SubstateProvider, {substateName: 'reflex___state____state', "
+        "contextName: 'reflex___state____state'}," in rendered
+    )
+    assert (
+        "createElement(SubstateProvider, {substateName: 'reflex___state____state__sub', "
+        "contextName: 'reflex___state____state__sub'}," in rendered
+    )
+    # The reducers moved into SubstateProvider; StateProvider only composes.
+    provider_body = rendered[rendered.index("export function StateProvider") :]
+    assert "useReducer" not in provider_body
+    assert "createElement(DispatchProvider, {}," in provider_body
+
+
 def test_context_template_client_side_component_is_named():
     """``ClientSide`` returns a named component, not an anonymous arrow."""
     from reflex_base.compiler.templates import context_template