Dynamic selectors

package.json

@@ -29,5 +29,8 @@
   "license": "MIT",
   "devDependencies": {
     "prettier": "^3.6.2"
+  },
+  "prettier": {
+    "trailingComma": "all"
   }
 }

packages/use-store/src/experimental/useStore.spec.tsx

@@ -998,58 +998,6 @@ describe("Experimental Userland Store", () => {
     expect(store._listeners.length).toBe(0);
   });
 
-  it("dynamic selectors are not yet supported", async () => {
-    const store = createStore(reducer, 1);
-
-    let setSelector: any;
-    function Count({ testid }: { testid: string }) {
-      const [selector, _setSelector] = useState(() => identity);
-      setSelector = _setSelector;
-      const count = useStoreSelector(store, selector);
-      logger.log({ testid, count });
-      return <div>{count}</div>;
-    }
-
-    function App() {
-      return (
-        <StoreProvider>
-          <Count testid="count" />
-        </StoreProvider>
-      );
-    }
-
-    const { asFragment, unmount } = await act(async () => {
-      return render(<App />);
-    });
-
-    logger.assertLog([{ testid: "count", count: 1 }]);
-
-    expect(asFragment()).toMatchInlineSnapshot(`
-      <DocumentFragment>
-        <div>
-          1
-        </div>
-      </DocumentFragment>
-    `);
-
-    let error: any;
-    try {
-      await act(async () => {
-        setSelector((s: number) => s * 2);
-      });
-    } catch (e) {
-      error = e;
-    }
-
-    logger.assertLog([]);
-
-    expect(error.message).toMatch(
-      "useStoreSelector does not currently support dynamic selectors",
-    );
-    unmount();
-    expect(store._listeners.length).toBe(0);
-  });
-
   it("dynamic stores are not yet supported", async () => {
     const store1 = createStore(reducer, 1);
     const store2 = createStore(reducer, 10);
@@ -1367,3 +1315,144 @@ describe("Experimental Userland Store", () => {
     expect(storeB._listeners.length).toBe(0);
   });
 });
+
+describe("Selectors can be dynamic", () => {
+  it("dynamic selectors are supported", async () => {
+    const store = createStore(reducer, 1);
+
+    let setSelector: any;
+    function Count({ testid }: { testid: string }) {
+      const [selector, _setSelector] = useState(() => identity);
+      setSelector = _setSelector;
+      const count = useStoreSelector(store, selector);
+      logger.log({ testid, count });
+      return <div>{count}</div>;
+    }
+
+    function App() {
+      return (
+        <StoreProvider>
+          <Count testid="count" />
+        </StoreProvider>
+      );
+    }
+
+    const { asFragment, unmount } = await act(async () => {
+      return render(<App />);
+    });
+
+    logger.assertLog([{ testid: "count", count: 1 }]);
+
+    expect(asFragment()).toMatchInlineSnapshot(`
+      <DocumentFragment>
+        <div>
+          1
+        </div>
+      </DocumentFragment>
+    `);
+
+    await act(async () => {
+      setSelector(() => (s: number) => s * 2);
+    });
+
+    logger.assertLog([{ testid: "count", count: 2 }]);
+
+    expect(asFragment()).toMatchInlineSnapshot(`
+      <DocumentFragment>
+        <div>
+          2
+        </div>
+      </DocumentFragment>
+    `);
+
+    unmount();
+    expect(store._listeners.length).toBe(0);
+  });
+
+  it("selector changes sync during a transition update to the store", async () => {
+    const store = createStore(reducer, 1);
+
+    let setSelector: any;
+    function Count({ testid }: { testid: string }) {
+      const [selector, _setSelector] = useState(() => identity);
+      setSelector = _setSelector;
+      const count = useStoreSelector(store, selector);
+      logger.log({ testid, count });
+      return <div>{count}</div>;
+    }
+
+    function App() {
+      return (
+        <StoreProvider>
+          <Count testid="count" />
+        </StoreProvider>
+      );
+    }
+
+    const { asFragment, unmount } = await act(async () => {
+      return render(<App />);
+    });
+
+    logger.assertLog([{ testid: "count", count: 1 }]);
+
+    expect(asFragment()).toMatchInlineSnapshot(`
+      <DocumentFragment>
+        <div>
+          1
+        </div>
+      </DocumentFragment>
+    `);
+
+    let resolve: () => void;
+
+    const promise = new Promise<void>((_resolve) => {
+      resolve = _resolve;
+    });
+
+    await act(async () => {
+      startTransition(async () => {
+        store.dispatch({ type: "INCREMENT" });
+        await promise;
+      });
+    });
+
+    await act(async () => {
+      setSelector(() => (s: number) => s * 2);
+    });
+
+    logger.assertLog([
+      // Just like a fresh mount, the new selector is run on the transition state...
+      { testid: "count", count: 4 },
+      // But gets fixed up to the sync state in the layoutEffect before yielding to the user.
+      { testid: "count", count: 2 },
+    ]);
+
+    expect(asFragment()).toMatchInlineSnapshot(`
+      <DocumentFragment>
+        <div>
+          2
+        </div>
+      </DocumentFragment>
+    `);
+
+    await act(async () => {
+      resolve();
+    });
+
+    logger.assertLog([
+      // Now we render the transition state
+      { testid: "count", count: 4 },
+    ]);
+
+    expect(asFragment()).toMatchInlineSnapshot(`
+      <DocumentFragment>
+        <div>
+          4
+        </div>
+      </DocumentFragment>
+    `);
+
+    unmount();
+    expect(store._listeners.length).toBe(0);
+  });
+});

packages/use-store/src/experimental/useStore.tsx

@@ -101,6 +101,11 @@ export function StoreProvider({ children }: { children: React.ReactNode }) {
   );
 }
 
+type HookState<S, T> = {
+  value: T;
+  selector: (state: S) => T;
+};
+
 /**
  * Tearing-resistant hook for consuming application state locally within a
  * component (without prop drilling or putting state in context).
@@ -147,12 +152,6 @@ export function useStoreSelector<S, T>(
       "useStoreSelector does not currently support dynamic stores",
     );
   }
-  const previousSelectorRef = useRef(selector);
-  if (selector !== previousSelectorRef.current) {
-    throw new Error(
-      "useStoreSelector does not currently support dynamic selectors",
-    );
-  }
 
   // Counterintuitively we initially render with the transition/head state
   // instead of the committed state. This is required in order for us to
@@ -166,14 +165,38 @@ export function useStoreSelector<S, T>(
   // Instead we must initially render with the transition state and then
   // trigger a sync fixup setState in the useLayoutEffect if we are mounting
   // sync and thus should be showing the committed state.
-  const [state, setState] = useState<T>(() => selector(store.getState()));
+  //
+  // We also track the selector used for each state so that we can determine if
+  // the selector has changed since our last updated.
+  const [hookState, setState] = useState<HookState<S, T>>(() => ({
+    value: selector(store.getState()),
+    selector,
+  }));
+
+  // If we have a new selector, we try to derive a new value during render. If
+  // the mount was sync, we'll apply a fixup in useLayoutEffect, just like we do
+  // on mount.
+  const selectorChange = hookState.selector !== selector;
+  const state = selectorChange ? selector(store.getState()) : hookState.value;
 
   useLayoutEffect(() => {
     // Ensure our store is managed by the tracker.
     storeManager.addStore(store);
     const mountState = selector(store.getState());
     const mountCommittedState = selector(store.getCommittedState());
 
+    // Helper to ensure we preserve object identity if neither state nor selector has changed.
+    function setHookState(value: T) {
+      setState((prev) => {
+        // If nothing has changed...
+        if (prev.value === value && prev.selector === selector) {
+          // Preserve object identity.
+          return prev;
+        }
+        return { value, selector };
+      });
+    }
+
     // If we are mounting as part of a sync update mid transition, our initial
     // render value was wrong and we must trigger a sync fixup update.
     // Similarly, if a sync state update was triggered between the moment we
@@ -183,7 +206,7 @@ export function useStoreSelector<S, T>(
     // Both of these cases manifest as our initial render state not matching
     // the currently committed state.
     if (state !== mountCommittedState) {
-      setState(mountCommittedState);
+      setHookState(mountCommittedState);
     }
 
     // If we mounted mid-transition, and that transition is still ongoing, we
@@ -200,20 +223,21 @@ export function useStoreSelector<S, T>(
       // while we were mounting resolves, it will also include rerendering
       // this component to reflect the new state.
       startTransition(() => {
-        setState(mountState);
+        setHookState(mountState);
       });
     }
+
     const unsubscribe = store.subscribe(() => {
-      const state = store.getState();
-      setState(selector(state));
+      setHookState(selector(store.getState()));
     });
     return () => {
       unsubscribe();
       storeManager.removeStore(store);
     };
-    // We intentionally ignore `state` since we only care about its value on mount
+    // We intentionally ignore `state` since we only care about its value on
+    // mount or when the selector changes.
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, []);
+  }, [selector]);
 
   return state;
 }